Skip to content

Commit

Permalink
Merge pull request #17 from ctdunc/chore/1.1_release
Browse files Browse the repository at this point in the history 
Upgrade compatibility with polars 1.*
  • Loading branch information
@abstractqqq
abstractqqq authored Jul 22, 2024
2 parents b20be36 + 69cd063 commit 2403f3e
Show file tree
Hide file tree
Showing 7 changed files with 139 additions and 436 deletions.
2 changes: 1 addition & 1 deletion pyproject.toml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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 = [
Expand Down
45 changes: 33 additions & 12 deletions python/polars_istr/_utils.py
View file
Original file line number Diff line number Diff line change
@@ -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):
Expand All @@ -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,
Expand All @@ -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,
Expand Down
133 changes: 17 additions & 116 deletions python/polars_istr/cusip.py
View file
Original file line number Diff line number Diff line change
@@ -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.
Expand All @@ -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,
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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)
Loading

0 comments on commit 2403f3e

Please sign in to comment.