Merge pull request #17 from ctdunc/chore/1.1_release

Upgrade compatibility with polars 1.*

Author: abstractqqq

pyproject.toml

@@ -5,7 +5,7 @@ build-backend = "maturin"
 [project]
 name = "polars_istr"
 requires-python = ">=3.8"
-version = "0.1.0"
+version = "0.1.1"
 
 license = {file = "LICENSE.txt"}
 classifiers = [

python/polars_istr/_utils.py

@@ -1,7 +1,31 @@
+import os
+import re
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Union
+
 import polars as pl
-from typing import Any, Optional, List, Dict
+
 from .type_alias import StrOrExpr
 
+_POLARS_LEGACY_SUPPORT = tuple(int(re.sub("[^0-9]", "", x)) for x in pl.__version__.split(".")) < (
+    0,
+    20,
+    16,
+)
+_IS_POLARS_V1 = pl.__version__.startswith("1")
+
+_PLUGIN_PATH = Path(__file__).parent
+
+_PLUGIN_LIB_LEGACY = os.path.join(
+    os.path.dirname(__file__),
+    next(
+        filter(
+            lambda file: file.endswith((".so", ".dll", ".pyd")),
+            os.listdir(os.path.dirname(__file__)),
+        )
+    ),
+)
+
 
 def str_to_expr(x: StrOrExpr) -> pl.Expr:
     if isinstance(x, str):
@@ -14,23 +38,20 @@ def str_to_expr(x: StrOrExpr) -> pl.Expr:
 
 def pl_plugin(
     *,
-    lib: str,
     symbol: str,
-    args: List[StrOrExpr],
+    args: List[Union[pl.Series, pl.Expr]],
     kwargs: Optional[Dict[str, Any]] = None,
     is_elementwise: bool = False,
     returns_scalar: bool = False,
     changes_length: bool = False,
     cast_to_supertype: bool = False,
 ) -> pl.Expr:
-    # pl.__version__ should always be a valid version number, so split returns always 3 strs
-    if tuple(int(x) for x in pl.__version__.split(".")) < (0, 20, 16):
-        # This will eventually be deprecated?
-        first = str_to_expr(args[0])
-        return first.register_plugin(
-            lib=lib,
+    if _POLARS_LEGACY_SUPPORT:
+        # This will eventually be deprecated, yes
+        return args[0].register_plugin(
+            lib=_PLUGIN_LIB_LEGACY,
             symbol=symbol,
-            args=[str_to_expr(x) for x in args[1:]],
+            args=args[1:],
             kwargs=kwargs,
             is_elementwise=is_elementwise,
             returns_scalar=returns_scalar,
@@ -41,8 +62,8 @@ def pl_plugin(
     from polars.plugins import register_plugin_function
 
     return register_plugin_function(
-        plugin_path=lib,
-        args=[str_to_expr(x) for x in args],
+        plugin_path=_PLUGIN_PATH,
+        args=args,
         function_name=symbol,
         kwargs=kwargs,
         is_elementwise=is_elementwise,

python/polars_istr/cusip.py

@@ -1,12 +1,9 @@
 from __future__ import annotations
 import polars as pl
-from polars.utils.udfs import _get_shared_lib_location
-from ._utils import pl_plugin, StrOrExpr
+from ._utils import pl_plugin
 
-_lib = _get_shared_lib_location(__file__)
 
-
-def cusip_extract_all(x: StrOrExpr) -> pl.Expr:
+def cusip_extract_all(x: pl.Series | pl.Expr) -> pl.Expr:
     """
     Returns a struct containing country_code, issue_num, issuer_num, check_digit,
     or null, if it cannot be parsed.
@@ -15,106 +12,98 @@ def cusip_extract_all(x: StrOrExpr) -> pl.Expr:
     """
     return pl_plugin(
         args=[x],
-        lib=_lib,
         symbol="pl_cusip_full",
         is_elementwise=True,
     )
 
 
-def cusip_issue_num(x: StrOrExpr) -> pl.Expr:
+def cusip_issue_num(x: pl.Series | pl.Expr) -> pl.Expr:
     """
     Returns the issue number from the CUSIP, or null if it cannot be parsed.
     """
     return pl_plugin(
         args=[x],
-        lib=_lib,
         symbol="pl_cusip_issue_num",
         is_elementwise=True,
     )
 
 
-def cusip_issuer_num(x: StrOrExpr) -> pl.Expr:
+def cusip_issuer_num(x: pl.Series | pl.Expr) -> pl.Expr:
     """
     Returns the issuer number from the CUSIP, or null if it cannot be parsed.
     """
     return pl_plugin(
         args=[x],
-        lib=_lib,
         symbol="pl_cusip_issuer_num",
         is_elementwise=True,
     )
 
 
-def cusip_check_digit(x: StrOrExpr) -> pl.Expr:
+def cusip_check_digit(x: pl.Series | pl.Expr) -> pl.Expr:
     """
     Returns check digit from the CUSIP, or null if it cannot be parsed.
     """
     return pl_plugin(
         args=[x],
-        lib=_lib,
         symbol="pl_cusip_check_digit",
         is_elementwise=True,
     )
 
 
-def cusip_country_code(x: StrOrExpr) -> pl.Expr:
+def cusip_country_code(x: pl.Series | pl.Expr) -> pl.Expr:
     """
     Returns the country code from the CUSIP, or null if it cannot be parsed.
     """
     return pl_plugin(
         args=[x],
-        lib=_lib,
         symbol="pl_cusip_country_code",
         is_elementwise=True,
     )
 
 
-def cusip_payload(x: StrOrExpr) -> pl.Expr:
+def cusip_payload(x: pl.Series | pl.Expr) -> pl.Expr:
     """
     Returns the payload (CUSIP ex. check digit) from the CUSIP, or null if it
     cannot be parsed.
     """
     return pl_plugin(
         args=[x],
-        lib=_lib,
         symbol="pl_cusip_payload",
         is_elementwise=True,
     )
 
 
-def cusip_is_private_issue(x: StrOrExpr) -> pl.Expr:
+def cusip_is_private_issue(x: pl.Series | pl.Expr) -> pl.Expr:
     """
     Returns true if the issue number is reserved for private use.
     """
     return pl_plugin(
         args=[x],
-        lib=_lib,
         symbol="pl_cusip_is_private_issue",
         is_elementwise=True,
     )
 
 
-def cusip_has_private_issuer(x: StrOrExpr) -> pl.Expr:
+def cusip_has_private_issuer(x: pl.Series | pl.Expr) -> pl.Expr:
     """
     Returns true if the issuer is reserved for private use.
     """
     return pl_plugin(
         args=[x],
-        lib=_lib,
         symbol="pl_cusip_has_private_issuer",
         is_elementwise=True,
     )
 
 
-def cusip_is_private_use(x: StrOrExpr) -> pl.Expr:
+def cusip_is_private_use(x: pl.Series | pl.Expr) -> pl.Expr:
     """
     Returns True if either the issuer or issue number is reserved for
     private use.
     """
-    return pl_plugin(args=[x], lib=_lib, symbol="pl_cusip_is_private_use", is_elementwise=True)
+    return pl_plugin(args=[x], symbol="pl_cusip_is_private_use", is_elementwise=True)
 
 
-def cusip_is_cins(x: StrOrExpr) -> pl.Expr:
+def cusip_is_cins(x: pl.Series | pl.Expr) -> pl.Expr:
     """
     Returns true if this CUSIP number is actually a
     CUSIP International Numbering System (CINS) number,
@@ -123,10 +112,10 @@ def cusip_is_cins(x: StrOrExpr) -> pl.Expr:
 
     Null if unable to parse.
     """
-    return pl_plugin(args=[x], lib=_lib, symbol="pl_cusip_is_cins", is_elementwise=True)
+    return pl_plugin(args=[x], symbol="pl_cusip_is_cins", is_elementwise=True)
 
 
-def cusip_is_cins_base(x: StrOrExpr) -> pl.Expr:
+def cusip_is_cins_base(x: pl.Series | pl.Expr) -> pl.Expr:
     """
     Returns true if this CUSIP identifier is actually a CUSIP International
     Numbering System (CINS) identifier (with the further restriction that
@@ -135,10 +124,10 @@ def cusip_is_cins_base(x: StrOrExpr) -> pl.Expr:
 
     Null if unable to parse.
     """
-    return pl_plugin(args=[x], lib=_lib, symbol="pl_cusip_is_cins_base", is_elementwise=True)
+    return pl_plugin(args=[x], symbol="pl_cusip_is_cins_base", is_elementwise=True)
 
 
-def cusip_is_cins_extended(x: StrOrExpr) -> pl.Expr:
+def cusip_is_cins_extended(x: pl.Series | pl.Expr) -> pl.Expr:
     """
     Returns true if this CUSIP identifier is actually a CUSIP International
     Numbering System (CINS) identifier (with the further restriction that
@@ -147,92 +136,4 @@ def cusip_is_cins_extended(x: StrOrExpr) -> pl.Expr:
 
     Null if unable to parse.
     """
-    return pl_plugin(args=[x], lib=_lib, symbol="pl_cusip_is_cins_extended", is_elementwise=True)
-
-
-@pl.api.register_expr_namespace("cusip")
-class CusipExt:
-    """
-    This class contains tools for parsing CUSIP/CINS and Extended CINS format data.
-
-    Polars Namespace: cusip
-
-    Example: pl.col("cusip_cins_string").cusip.country_code()
-    """
-
-    def __init__(self, expr: pl.Expr):
-        self._expr: pl.Expr = expr
-
-    def extract_all(self) -> pl.Expr:
-        return self._expr.register_plugin(
-            lib=_lib,
-            symbol="pl_cusip_full",
-            is_elementwise=True,
-        )
-
-    def issue_num(self) -> pl.Expr:
-        return self._expr.register_plugin(
-            lib=_lib,
-            symbol="pl_cusip_issue_num",
-            is_elementwise=True,
-        )
-
-    def issuer_num(self) -> pl.Expr:
-        return self._expr.register_plugin(
-            lib=_lib,
-            symbol="pl_cusip_issuer_num",
-            is_elementwise=True,
-        )
-
-    def check_digit(self) -> pl.Expr:
-        return self._expr.register_plugin(
-            lib=_lib,
-            symbol="pl_cusip_check_digit",
-            is_elementwise=True,
-        )
-
-    def country_code(self) -> pl.Expr:
-        return self._expr.register_plugin(
-            lib=_lib,
-            symbol="pl_cusip_country_code",
-            is_elementwise=True,
-        )
-
-    def payload(self) -> pl.Expr:
-        return self._expr.register_plugin(
-            lib=_lib,
-            symbol="pl_cusip_payload",
-            is_elementwise=True,
-        )
-
-    def is_private_issue(self) -> pl.Expr:
-        return self._expr.register_plugin(
-            lib=_lib,
-            symbol="pl_cusip_is_private_issue",
-            is_elementwise=True,
-        )
-
-    def has_private_issuer(self) -> pl.Expr:
-        return self._expr.register_plugin(
-            lib=_lib,
-            symbol="pl_cusip_has_private_issuer",
-            is_elementwise=True,
-        )
-
-    def is_private_use(self) -> pl.Expr:
-        return self._expr.register_plugin(
-            lib=_lib, symbol="pl_cusip_is_private_use", is_elementwise=True
-        )
-
-    def is_cins(self) -> pl.Expr:
-        return self._expr.register_plugin(lib=_lib, symbol="pl_cusip_is_cins", is_elementwise=True)
-
-    def is_cins_base(self) -> pl.Expr:
-        return self._expr.register_plugin(
-            lib=_lib, symbol="pl_cusip_is_cins_base", is_elementwise=True
-        )
-
-    def is_cins_extended(self) -> pl.Expr:
-        return self._expr.register_plugin(
-            lib=_lib, symbol="pl_cusip_is_cins_extended", is_elementwise=True
-        )
+    return pl_plugin(args=[x], symbol="pl_cusip_is_cins_extended", is_elementwise=True)