fix: Modernize expression_lib (#111)

Author: MarcoGorelli

README.md

@@ -56,20 +56,24 @@ fn pig_latinnify(inputs: &[Series], kwargs: PigLatinKwargs) -> PolarsResult<Seri
 This can then be exposed on the Python side:
 
 ```python
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
 import polars as pl
-from polars.type_aliases import IntoExpr
-from polars.utils.udfs import _get_shared_lib_location
+from polars.plugins import register_plugin_function
 
-from expression_lib.utils import parse_into_expr
+from expression_lib._utils import LIB
 
-lib = _get_shared_lib_location(__file__)
+if TYPE_CHECKING:
+    from expression_lib._typing import IntoExprColumn
 
 
-def pig_latinnify(expr: IntoExpr, capitalize: bool = False) -> pl.Expr:
-    expr = parse_into_expr(expr)
-    return expr.register_plugin(
-        lib=lib,
-        symbol="pig_latinnify",
+def pig_latinnify(expr: IntoExprColumn, capitalize: bool = False) -> pl.Expr:
+    return register_plugin_function(
+        plugin_path=LIB,
+        args=[expr],
+        function_name="pig_latinnify",
         is_elementwise=True,
         kwargs={"capitalize": capitalize},
     )

example/derive_expression/expression_lib/expression_lib/_typing.py

@@ -0,0 +1,13 @@
+from typing import TYPE_CHECKING, Union
+
+if TYPE_CHECKING:
+    import sys
+
+    import polars as pl
+
+    if sys.version_info >= (3, 10):
+        from typing import TypeAlias
+    else:
+        from typing_extensions import TypeAlias
+
+    IntoExprColumn: TypeAlias = Union[pl.Expr, str, pl.Series]

example/derive_expression/expression_lib/expression_lib/_utils.py

@@ -0,0 +1,3 @@
+from pathlib import Path
+
+LIB = Path(__file__).parent

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 IntoExprColumn
 
-def is_leap_year(expr: IntoExpr) -> pl.Expr:
-    expr = parse_into_expr(expr)
+
+def is_leap_year(expr: IntoExprColumn) -> pl.Expr:
     return register_plugin_function(
-        plugin_path=Path(__file__).parent,
+        plugin_path=LIB,
         args=[expr],
         function_name="is_leap_year",
         is_elementwise=True,
@@ -18,10 +22,11 @@ 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)
+def change_time_zone(expr: IntoExprColumn, tz: str = "Europe/Amsterdam") -> pl.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},
     )

example/derive_expression/expression_lib/expression_lib/dist.py

@@ -1,40 +1,42 @@
+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 IntoExprColumn
 
-from expression_lib.utils import parse_into_expr
 
-def hamming_distance(expr: IntoExpr, other: IntoExpr) -> pl.Expr:
-    expr = parse_into_expr(expr)
+def hamming_distance(expr: IntoExprColumn, other: IntoExprColumn) -> pl.Expr:
     return register_plugin_function(
-        plugin_path=Path(__file__).parent,
+        plugin_path=LIB,
         args=[expr, other],
         function_name="hamming_distance",
         is_elementwise=True,
     )
 
 
-def jaccard_similarity(expr: IntoExpr, other: IntoExpr) -> pl.Expr:
-    expr = parse_into_expr(expr)
+def jaccard_similarity(expr: IntoExprColumn, other: IntoExprColumn) -> pl.Expr:
     return register_plugin_function(
-        plugin_path=Path(__file__).parent,
+        plugin_path=LIB,
         args=[expr, other],
         function_name="jaccard_similarity",
         is_elementwise=True,
     )
 
 
 def haversine(
-    start_lat: IntoExpr,
-    start_long: IntoExpr,
-    end_lat: IntoExpr,
-    end_long: IntoExpr,
+    start_lat: IntoExprColumn,
+    start_long: IntoExprColumn,
+    end_lat: IntoExprColumn,
+    end_long: IntoExprColumn,
 ) -> 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,

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")

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 IntoExprColumn
 
 
-def pig_latinnify(expr: IntoExpr, capitalize: bool = False) -> pl.Expr:
-    expr = parse_into_expr(expr)
+def pig_latinnify(expr: IntoExprColumn, capitalize: bool = False) -> pl.Expr:
     return register_plugin_function(
-        plugin_path=Path(__file__).parent,
+        plugin_path=LIB,
         args=[expr],
         function_name="pig_latinnify",
         is_elementwise=True,
@@ -19,7 +22,7 @@ def pig_latinnify(expr: IntoExpr, capitalize: bool = False) -> pl.Expr:
 
 
 def append_args(
-    expr: IntoExpr,
+    expr: IntoExprColumn,
     float_arg: float,
     integer_arg: int,
     string_arg: str,
@@ -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,

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 IntoExprColumn
 
 
-def panic(expr: IntoExpr) -> pl.Expr:
-    expr = parse_into_expr(expr)
+def panic(expr: IntoExprColumn) -> pl.Expr:
     return register_plugin_function(
-        plugin_path=Path(__file__).parent,
+        plugin_path=LIB,
         args=[expr],
         function_name="panic",
     )