From f87f5b8c23fc5bb0c4b58ba963a008634e5055be Mon Sep 17 00:00:00 2001 From: Marco Gorelli <33491632+MarcoGorelli@users.noreply.github.com> Date: Wed, 2 Oct 2024 09:50:09 +0100 Subject: [PATCH] modernize expression_lib --- README.md | 19 ++++---- .../expression_lib/date_util.py | 21 ++++---- .../expression_lib/expression_lib/dist.py | 20 ++++---- .../expression_lib/extension.py | 7 ++- .../expression_lib/expression_lib/language.py | 16 ++++--- .../expression_lib/expression_lib/panic.py | 14 ++++-- .../expression_lib/expression_lib/utils.py | 48 ------------------- 7 files changed, 57 insertions(+), 88 deletions(-) delete mode 100644 example/derive_expression/expression_lib/expression_lib/utils.py diff --git a/README.md b/README.md index a8b66c7..cb98aa7 100644 --- a/README.md +++ b/README.md @@ -56,20 +56,21 @@ fn pig_latinnify(inputs: &[Series], kwargs: PigLatinKwargs) -> PolarsResult pl.Expr: - expr = parse_into_expr(expr) - return expr.register_plugin( - lib=lib, - symbol="pig_latinnify", + return register_plugin_function( + plugin_path=LIB, + args=[expr], + function_name="pig_latinnify", is_elementwise=True, kwargs={"capitalize": capitalize}, ) diff --git a/example/derive_expression/expression_lib/expression_lib/date_util.py b/example/derive_expression/expression_lib/expression_lib/date_util.py index 928d7b8..14e03af 100644 --- a/example/derive_expression/expression_lib/expression_lib/date_util.py +++ b/example/derive_expression/expression_lib/expression_lib/date_util.py @@ -1,15 +1,19 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + import polars as pl -from polars.type_aliases import IntoExpr from polars.plugins import register_plugin_function -from pathlib import Path -from expression_lib.utils import parse_into_expr +from expression_lib._utils import LIB + +if TYPE_CHECKING: + from expression_lib._typing import IntoExpr def is_leap_year(expr: IntoExpr) -> pl.Expr: - expr = parse_into_expr(expr) return register_plugin_function( - plugin_path=Path(__file__).parent, + plugin_path=LIB, args=[expr], function_name="is_leap_year", is_elementwise=True, @@ -19,9 +23,10 @@ def is_leap_year(expr: IntoExpr) -> pl.Expr: # Note that this already exists in Polars. It is just for explanatory # purposes. def change_time_zone(expr: IntoExpr, tz: str = "Europe/Amsterdam") -> pl.Expr: - expr = parse_into_expr(expr) return register_plugin_function( - plugin_path=Path(__file__).parent, + plugin_path=LIB, args=[expr], - function_name="change_time_zone", is_elementwise=True, kwargs={"tz": tz} + function_name="change_time_zone", + is_elementwise=True, + kwargs={"tz": tz}, ) diff --git a/example/derive_expression/expression_lib/expression_lib/dist.py b/example/derive_expression/expression_lib/expression_lib/dist.py index e9274e6..1d40f7a 100644 --- a/example/derive_expression/expression_lib/expression_lib/dist.py +++ b/example/derive_expression/expression_lib/expression_lib/dist.py @@ -1,15 +1,19 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + import polars as pl -from polars.type_aliases import IntoExpr from polars.plugins import register_plugin_function -from pathlib import Path +from expression_lib._utils import LIB + +if TYPE_CHECKING: + from expression_lib._typing import IntoExpr -from expression_lib.utils import parse_into_expr def hamming_distance(expr: IntoExpr, other: IntoExpr) -> pl.Expr: - expr = parse_into_expr(expr) return register_plugin_function( - plugin_path=Path(__file__).parent, + plugin_path=LIB, args=[expr, other], function_name="hamming_distance", is_elementwise=True, @@ -17,9 +21,8 @@ def hamming_distance(expr: IntoExpr, other: IntoExpr) -> pl.Expr: def jaccard_similarity(expr: IntoExpr, other: IntoExpr) -> pl.Expr: - expr = parse_into_expr(expr) return register_plugin_function( - plugin_path=Path(__file__).parent, + plugin_path=LIB, args=[expr, other], function_name="jaccard_similarity", is_elementwise=True, @@ -32,9 +35,8 @@ def haversine( end_lat: IntoExpr, end_long: IntoExpr, ) -> pl.Expr: - start_lat = parse_into_expr(start_lat) return register_plugin_function( - plugin_path=Path(__file__).parent, + plugin_path=LIB, args=[start_lat, start_long, end_lat, end_long], function_name="haversine", is_elementwise=True, diff --git a/example/derive_expression/expression_lib/expression_lib/extension.py b/example/derive_expression/expression_lib/expression_lib/extension.py index c286424..fbd9bc5 100644 --- a/example/derive_expression/expression_lib/expression_lib/extension.py +++ b/example/derive_expression/expression_lib/expression_lib/extension.py @@ -16,11 +16,14 @@ - static typing will not recognise your custom namespace. Errors such as `"Expr" has no attribute "dist" [attr-defined]`. """ + from __future__ import annotations -import polars as pl from typing import Any, Callable -from expression_lib import date_util, dist, language, utils, panic + +import polars as pl + +from expression_lib import date_util, dist, language, panic @pl.api.register_expr_namespace("language") diff --git a/example/derive_expression/expression_lib/expression_lib/language.py b/example/derive_expression/expression_lib/expression_lib/language.py index 0713f0d..1c9c0c5 100644 --- a/example/derive_expression/expression_lib/expression_lib/language.py +++ b/example/derive_expression/expression_lib/expression_lib/language.py @@ -1,16 +1,19 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + import polars as pl -from polars.type_aliases import IntoExpr from polars.plugins import register_plugin_function -from pathlib import Path -from expression_lib.utils import parse_into_expr +from expression_lib._utils import LIB +if TYPE_CHECKING: + from expression_lib._typing import IntoExpr def pig_latinnify(expr: IntoExpr, capitalize: bool = False) -> pl.Expr: - expr = parse_into_expr(expr) return register_plugin_function( - plugin_path=Path(__file__).parent, + plugin_path=LIB, args=[expr], function_name="pig_latinnify", is_elementwise=True, @@ -28,9 +31,8 @@ def append_args( """ This example shows how arguments other than `Series` can be used. """ - expr = parse_into_expr(expr) return register_plugin_function( - plugin_path=Path(__file__).parent, + plugin_path=LIB, args=[expr], kwargs={ "float_arg": float_arg, diff --git a/example/derive_expression/expression_lib/expression_lib/panic.py b/example/derive_expression/expression_lib/expression_lib/panic.py index 6338a76..297ede7 100644 --- a/example/derive_expression/expression_lib/expression_lib/panic.py +++ b/example/derive_expression/expression_lib/expression_lib/panic.py @@ -1,15 +1,19 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + import polars as pl -from polars.type_aliases import IntoExpr from polars.plugins import register_plugin_function -from expression_lib.utils import parse_into_expr -from pathlib import Path +from expression_lib._utils import LIB + +if TYPE_CHECKING: + from expression_lib._typing import IntoExpr def panic(expr: IntoExpr) -> pl.Expr: - expr = parse_into_expr(expr) return register_plugin_function( - plugin_path=Path(__file__).parent, + plugin_path=LIB, args=[expr], function_name="panic", ) diff --git a/example/derive_expression/expression_lib/expression_lib/utils.py b/example/derive_expression/expression_lib/expression_lib/utils.py deleted file mode 100644 index 73f7a19..0000000 --- a/example/derive_expression/expression_lib/expression_lib/utils.py +++ /dev/null @@ -1,48 +0,0 @@ -from __future__ import annotations - -from typing import TYPE_CHECKING - -import polars as pl - -if TYPE_CHECKING: - from polars.type_aliases import IntoExpr, PolarsDataType - - -def parse_into_expr( - expr: IntoExpr, - *, - str_as_lit: bool = False, - list_as_lit: bool = True, - dtype: PolarsDataType | None = None, -) -> pl.Expr: - """ - Parse a single input into an expression. - - Parameters - ---------- - expr - The input to be parsed as an expression. - str_as_lit - Interpret string input as a string literal. If set to `False` (default), - strings are parsed as column names. - list_as_lit - Interpret list input as a lit literal, If set to `False`, - lists are parsed as `Series` literals. - dtype - If the input is expected to resolve to a literal with a known dtype, pass - this to the `lit` constructor. - - Returns - ------- - polars.Expr - """ - if isinstance(expr, pl.Expr): - pass - elif isinstance(expr, str) and not str_as_lit: - expr = pl.col(expr) - elif isinstance(expr, list) and not list_as_lit: - expr = pl.lit(pl.Series(expr), dtype=dtype) - else: - expr = pl.lit(expr, dtype=dtype) - - return expr