Merge pull request #5 from abstractqqq/update_example

update_example
abstractqqq · Mar 4, 2024 · 0f21e51 · 0f21e51
2 parents 35af446 + 9a0a2ad
commit 0f21e51
Show file tree

Hide file tree

Showing 2 changed files with 151 additions and 42 deletions.
diff --git a/examples/basics.ipynb b/examples/basics.ipynb
@@ -156,19 +156,121 @@
     ").unnest(\"ib\")"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "4a6b3620",
+   "metadata": {},
+   "source": [
+    "# ISIN"
+   ]
+  },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 5,
    "id": "787f6807",
    "metadata": {},
-   "outputs": [],
-   "source": []
+   "outputs": [
+    {
+     "data": {
+      "text/html": [
+       "<div><style>\n",
+       ".dataframe > thead > tr,\n",
+       ".dataframe > tbody > tr {\n",
+       "  text-align: right;\n",
+       "  white-space: pre-wrap;\n",
+       "}\n",
+       "</style>\n",
+       "<small>shape: (5, 1)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>isin</th></tr><tr><td>str</td></tr></thead><tbody><tr><td>&quot;US0378331005&quot;</td></tr><tr><td>&quot;US0378331008&quot;</td></tr><tr><td>&quot;US037833100&quot;</td></tr><tr><td>&quot;CA00206RGB20&quot;</td></tr><tr><td>&quot;XS1550212416&quot;</td></tr></tbody></table></div>"
+      ],
+      "text/plain": [
+       "shape: (5, 1)\n",
+       "┌──────────────┐\n",
+       "│ isin         │\n",
+       "│ ---          │\n",
+       "│ str          │\n",
+       "╞══════════════╡\n",
+       "│ US0378331005 │\n",
+       "│ US0378331008 │\n",
+       "│ US037833100  │\n",
+       "│ CA00206RGB20 │\n",
+       "│ XS1550212416 │\n",
+       "└──────────────┘"
+      ]
+     },
+     "execution_count": 5,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "df = pl.DataFrame(\n",
+    "    {\n",
+    "        \"isin\": [\n",
+    "            \"US0378331005\", # AAPL\n",
+    "            \"US0378331008\", # AAPL w/ bad check digit\n",
+    "            \"US037833100\", # AAPL w/o check digit\n",
+    "            \"CA00206RGB20\", # Canadian\n",
+    "            \"XS1550212416\", # Other\n",
+    "            None,\n",
+    "        ]\n",
+    "    }\n",
+    ")\n",
+    "df.head()"
+   ]
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 6,
    "id": "f4cb10f0",
    "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/html": [
+       "<div><style>\n",
+       ".dataframe > thead > tr,\n",
+       ".dataframe > tbody > tr {\n",
+       "  text-align: right;\n",
+       "  white-space: pre-wrap;\n",
+       "}\n",
+       "</style>\n",
+       "<small>shape: (6, 4)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>country_code</th><th>check_digit</th><th>security_id</th><th>is_valid</th></tr><tr><td>str</td><td>str</td><td>str</td><td>bool</td></tr></thead><tbody><tr><td>&quot;US&quot;</td><td>&quot;5&quot;</td><td>&quot;037833100&quot;</td><td>true</td></tr><tr><td>null</td><td>null</td><td>null</td><td>false</td></tr><tr><td>null</td><td>null</td><td>null</td><td>false</td></tr><tr><td>&quot;CA&quot;</td><td>&quot;0&quot;</td><td>&quot;00206RGB2&quot;</td><td>true</td></tr><tr><td>&quot;XS&quot;</td><td>&quot;6&quot;</td><td>&quot;155021241&quot;</td><td>true</td></tr><tr><td>null</td><td>null</td><td>null</td><td>false</td></tr></tbody></table></div>"
+      ],
+      "text/plain": [
+       "shape: (6, 4)\n",
+       "┌──────────────┬─────────────┬─────────────┬──────────┐\n",
+       "│ country_code ┆ check_digit ┆ security_id ┆ is_valid │\n",
+       "│ ---          ┆ ---         ┆ ---         ┆ ---      │\n",
+       "│ str          ┆ str         ┆ str         ┆ bool     │\n",
+       "╞══════════════╪═════════════╪═════════════╪══════════╡\n",
+       "│ US           ┆ 5           ┆ 037833100   ┆ true     │\n",
+       "│ null         ┆ null        ┆ null        ┆ false    │\n",
+       "│ null         ┆ null        ┆ null        ┆ false    │\n",
+       "│ CA           ┆ 0           ┆ 00206RGB2   ┆ true     │\n",
+       "│ XS           ┆ 6           ┆ 155021241   ┆ true     │\n",
+       "│ null         ┆ null        ┆ null        ┆ false    │\n",
+       "└──────────────┴─────────────┴─────────────┴──────────┘"
+      ]
+     },
+     "execution_count": 6,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "df.select(\n",
+    "    pl.col(\"isin\").isin.country_code().alias(\"country_code\"),\n",
+    "    pl.col(\"isin\").isin.check_digit().alias(\"check_digit\"),\n",
+    "    pl.col(\"isin\").isin.security_id().alias(\"security_id\"),\n",
+    "    pl.col(\"isin\").isin.is_valid().alias(\"is_valid\"),\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "9d5b6cbc",
+   "metadata": {},
    "outputs": [],
    "source": []
   }

diff --git a/tests/test_correctness.py b/tests/test_correctness.py
@@ -8,6 +8,7 @@
 from polars.testing import assert_frame_equal
 from typing import List
 
+
 @pytest.mark.parametrize(
     "df, cc, cd, reason, is_valid, bban, bank_id, branch_id",
     [
@@ -81,56 +82,62 @@ def test_iban1(
     assert_frame_equal(test1, ans)
     assert_frame_equal(test2, ans)
 
+
 @pytest.mark.parametrize(
-        "df, cc, cd, sec_id, is_valid",
-        [
-            (
-                pl.DataFrame(
-                    {
-                        "isin": [
-                            "US0378331005", # AAPL
-                            "US0378331008", # AAPL w/ bad check digit
-                            "US037833100", # AAPL w/o check digit
-                            "CA00206RGB20", # Canadian
-                            "XS1550212416", # Other
-                            None,
-                        ]
-                    }
-                ),
-            ["US", None, None, "CA", "XS",None],
+    "df, cc, cd, sec_id, is_valid",
+    [
+        (
+            pl.DataFrame(
+                {
+                    "isin": [
+                        "US0378331005",  # AAPL
+                        "US0378331008",  # AAPL w/ bad check digit
+                        "US037833100",  # AAPL w/o check digit
+                        "CA00206RGB20",  # Canadian
+                        "XS1550212416",  # Other
+                        None,
+                    ]
+                }
+            ),
+            ["US", None, None, "CA", "XS", None],
             ["5", None, None, "0", "6", None],
             ["037833100", None, None, "00206RGB2", "155021241", None],
-            [True, False, False, True, True, False]
-            )
-        ],
+            [True, False, False, True, True, False],
+        )
+    ],
 )
 def test_isin1(
-        df: pl.DataFrame,
-        cc: List[str],
-        cd: List[str],
-        sec_id: List[str],
-        is_valid: List[str],
-        ):
+    df: pl.DataFrame,
+    cc: List[str],
+    cd: List[str],
+    sec_id: List[str],
+    is_valid: List[str],
+):
     test1 = df.select(
+        pl.col("isin").isin.country_code().alias("country_code"),
+        pl.col("isin").isin.check_digit().alias("check_digit"),
+        pl.col("isin").isin.security_id().alias("security_id"),
+        pl.col("isin").isin.is_valid().alias("is_valid"),
+    )
+    test2 = (
+        df.lazy()
+        .select(
             pl.col("isin").isin.country_code().alias("country_code"),
             pl.col("isin").isin.check_digit().alias("check_digit"),
             pl.col("isin").isin.security_id().alias("security_id"),
             pl.col("isin").isin.is_valid().alias("is_valid"),
         )
-    test2 = df.lazy().select(
-            pl.col("isin").isin.country_code().alias("country_code"),
-            pl.col("isin").isin.check_digit().alias("check_digit"),
-            pl.col("isin").isin.security_id().alias("security_id"),
-            pl.col("isin").isin.is_valid().alias("is_valid"),
-        ).collect()
+        .collect()
+    )
 
-    ans = pl.DataFrame({
-        "country_code": cc,
-        "check_digit": cd,
-        "security_id": sec_id,
-        "is_valid": is_valid,
-    })
+    ans = pl.DataFrame(
+        {
+            "country_code": cc,
+            "check_digit": cd,
+            "security_id": sec_id,
+            "is_valid": is_valid,
+        }
+    )
 
     assert_frame_equal(test1, ans)
     assert_frame_equal(test2, ans)
-