Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Upgrade compatibility with polars 1.* #17

Merged
merged 2 commits into from
Jul 22, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion pyproject.toml
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
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
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
Loading