Flowfile 0.5.1__py3-none-any.whl → 0.5.3__py3-none-any.whl

flowfile_frame/lazy_methods.py

@@ -1,17 +1,33 @@
-import polars as pl
+from collections.abc import Callable
 from functools import wraps
-from typing import Optional, TypeVar, Type, Callable
-from flowfile_frame.utils import _get_function_source
+from typing import TypeVar
+
+import polars as pl
+
 from flowfile_frame.config import logger
+from flowfile_frame.utils import _get_function_source
 
-T = TypeVar('T')
-FlowFrameT = TypeVar('FlowFrameT', bound='FlowFrame')
+T = TypeVar("T")
+FlowFrameT = TypeVar("FlowFrameT", bound="FlowFrame")
 
 PASSTHROUGH_METHODS = {
-    'collect', 'collect_async', 'profile', 'describe', 'explain',
-    'show_graph', 'fetch', 'collect_schema', 'columns', 'dtypes',
-    'schema', 'width', 'estimated_size', 'n_chunks', 'is_empty',
-    'chunk_lengths', 'get_meta'
+    "collect",
+    "collect_async",
+    "profile",
+    "describe",
+    "explain",
+    "show_graph",
+    "fetch",
+    "collect_schema",
+    "columns",
+    "dtypes",
+    "schema",
+    "width",
+    "estimated_size",
+    "n_chunks",
+    "is_empty",
+    "chunk_lengths",
+    "get_meta",
 }
 
 
@@ -33,31 +49,68 @@ def create_lazyframe_method_wrapper(method_name: str, original_method: Callable)
     """
     # Determine if the original method returns a LazyFrame based on known method names
     lazyframe_returning_methods = {
-        "drop", "select", "with_columns", "sort", "filter", "join", "head", "tail",
-        "limit", "drop_nulls", "fill_null", "with_row_index", "group_by", "explode",
-        "unique", "slice", "shift", "reverse", "max", "min", "sum", "mean", "median",
-        "std", "var", "drop_nans", "fill_nan", "interpolate", "null_count", "quantile",
-        "unpivot", "melt", "first", "last"
+        "drop",
+        "select",
+        "with_columns",
+        "sort",
+        "filter",
+        "join",
+        "head",
+        "tail",
+        "limit",
+        "drop_nulls",
+        "fill_null",
+        "with_row_index",
+        "group_by",
+        "explode",
+        "unique",
+        "slice",
+        "shift",
+        "reverse",
+        "max",
+        "min",
+        "sum",
+        "mean",
+        "median",
+        "std",
+        "var",
+        "drop_nans",
+        "fill_nan",
+        "interpolate",
+        "null_count",
+        "quantile",
+        "unpivot",
+        "melt",
+        "first",
+        "last",
     }
 
     non_lazyframe_methods = {
-        "collect", "collect_schema", "fetch", "columns", "dtypes", "schema", "width",
-        "describe", "explain", "profile", "show_graph"
+        "collect",
+        "collect_schema",
+        "fetch",
+        "columns",
+        "dtypes",
+        "schema",
+        "width",
+        "describe",
+        "explain",
+        "profile",
+        "show_graph",
     }
 
-    returns_lazyframe = (
-            method_name in lazyframe_returning_methods or
-            (method_name not in non_lazyframe_methods and not method_name.startswith("_"))
+    returns_lazyframe = method_name in lazyframe_returning_methods or (
+        method_name not in non_lazyframe_methods and not method_name.startswith("_")
     )
 
     @wraps(original_method)
-    def wrapper(self, *args, description: Optional[str] = None, **kwargs):
+    def wrapper(self, *args, description: str | None = None, **kwargs):
         # Import here to avoid circular imports
         from flowfile_frame.flow_frame import generate_node_id
+
         new_node_id = generate_node_id()
 
-        if not all([True if not hasattr(arg, "convertable_to_code") else getattr(arg, 'convertable_to_code') for arg in
-                    args]):
+        if not all([True if not hasattr(arg, "convertable_to_code") else arg.convertable_to_code for arg in args]):
             logger.debug("Warning, could not create a good node")
             return self.__class__(getattr(self.data, method_name)(arg.expr for arg in args), flow_graph=self.flow_graph)
 
@@ -72,7 +125,7 @@ def create_lazyframe_method_wrapper(method_name: str, original_method: Callable)
                 # Try to get function source
                 try:
                     source, is_module_level = _get_function_source(arg)
-                    if source and hasattr(arg, '__name__') and arg.__name__ != '<lambda>':
+                    if source and hasattr(arg, "__name__") and arg.__name__ != "<lambda>":
                         function_sources.append(source)
                         # Use the function name in the representation
                         args_representations.append(arg.__name__)
@@ -89,7 +142,7 @@ def create_lazyframe_method_wrapper(method_name: str, original_method: Callable)
                 # Try to get function source
                 try:
                     source, is_module_level = _get_function_source(value)
-                    if source and hasattr(value, '__name__') and value.__name__ != '<lambda>':
+                    if source and hasattr(value, "__name__") and value.__name__ != "<lambda>":
                         function_sources.append(source)
                         kwargs_representations.append(f"{key}={value.__name__}")
                     else:
@@ -166,8 +219,7 @@ def add_lazyframe_methods(cls):
 
     # Skip properties and private methods
     skip_methods = {
-        name for name in dir(pl.LazyFrame)
-        if name.startswith('_') or isinstance(getattr(pl.LazyFrame, name), property)
+        name for name in dir(pl.LazyFrame) if name.startswith("_") or isinstance(getattr(pl.LazyFrame, name), property)
     }
 
     # Add all public LazyFrame methods that don't already exist
@@ -176,8 +228,8 @@ def add_lazyframe_methods(cls):
             continue
         attr = getattr(pl.LazyFrame, name)
         if name in PASSTHROUGH_METHODS:
-            def create_passthrough_method(method_name, method_attr):
 
+            def create_passthrough_method(method_name, method_attr):
                 @wraps(method_attr)
                 def passthrough_method(self, *args, **kwargs):
                     return getattr(self.data, method_name)(*args, **kwargs)
@@ -193,9 +245,10 @@ def add_lazyframe_methods(cls):
                 setattr(cls, name, wrapped_method)
 
     overlap = {
-        name for name in existing_methods
-        if name in dir(pl.LazyFrame) and not name.startswith('_') and callable(getattr(pl.LazyFrame, name))
+        name
+        for name in existing_methods
+        if name in dir(pl.LazyFrame) and not name.startswith("_") and callable(getattr(pl.LazyFrame, name))
     }
     if overlap:
         logger.debug(f"Preserved existing methods in {cls.__name__}: {', '.join(sorted(overlap))}")
-    return cls
+    return cls

flowfile_frame/list_name_space.py

@@ -1,27 +1,32 @@
 from __future__ import annotations
 
-from typing import Any, Optional, Union, TYPE_CHECKING, List, TypeVar, Callable, Sequence, Literal
+from collections.abc import Callable, Sequence
+from typing import TYPE_CHECKING, Any
 
 import polars as pl
 
-
 # --- TYPE CHECKING IMPORTS ---
 if TYPE_CHECKING:
-    from flowfile_frame.expr import Expr, _get_expr_and_repr, col, lit
-    from polars._typing import IntoExprColumn, NullBehavior, ListToStructWidthStrategy
     from datetime import date, datetime, time
 
+    from polars._typing import IntoExprColumn, ListToStructWidthStrategy, NullBehavior
+
+    from flowfile_frame.expr import Expr
+
 
 class ExprListNameSpace:
     """Namespace for list related expressions."""
 
-    def __init__(self, parent_expr: 'Expr', parent_repr_str: str):
+    def __init__(self, parent_expr: Expr, parent_repr_str: str):
         self.parent = parent_expr
         self.expr = parent_expr.expr.list if parent_expr.expr is not None else None
         self.parent_repr_str = parent_repr_str
 
-    def _create_next_expr(self, *args, method_name: str, result_expr: Optional[pl.Expr], is_complex: bool = True, **kwargs) -> 'Expr':
+    def _create_next_expr(
+        self, *args, method_name: str, result_expr: pl.Expr | None, is_complex: bool = True, **kwargs
+    ) -> Expr:
         from flowfile_frame.expr import Expr
+
         """Creates a new Expr instance, appending method call to repr string."""
         args_repr = ", ".join(repr(a) for a in args)
         kwargs_repr = ", ".join(f"{k}={repr(v)}" for k, v in kwargs.items())
@@ -46,7 +51,7 @@ class ExprListNameSpace:
             selector=None,
             agg_func=self.parent.agg_func,
             is_complex=is_complex,
-            convertable_to_code=self.parent.convertable_to_code
+            convertable_to_code=self.parent.convertable_to_code,
         )
         return new_expr_instance
 
@@ -67,13 +72,13 @@ class ExprListNameSpace:
         return self._create_next_expr(method_name="drop_nulls", result_expr=res_expr)
 
     def sample(
-            self,
-            n: int | IntoExprColumn | None = None,
-            *,
-            fraction: float | IntoExprColumn | None = None,
-            with_replacement: bool = False,
-            shuffle: bool = False,
-            seed: int | None = None,
+        self,
+        n: int | IntoExprColumn | None = None,
+        *,
+        fraction: float | IntoExprColumn | None = None,
+        with_replacement: bool = False,
+        shuffle: bool = False,
+        seed: int | None = None,
     ) -> Expr:
         if n is not None and fraction is not None:
             raise ValueError("cannot specify both `n` and `fraction`")
@@ -82,15 +87,15 @@ class ExprListNameSpace:
         if self.expr is not None:
             try:
                 if fraction is not None:
-                    expr_fraction = fraction.expr if hasattr(fraction, 'expr') else fraction
-                    res_expr = self.expr.sample(n=None, fraction=expr_fraction,
-                                                with_replacement=with_replacement,
-                                                shuffle=shuffle, seed=seed)
+                    expr_fraction = fraction.expr if hasattr(fraction, "expr") else fraction
+                    res_expr = self.expr.sample(
+                        n=None, fraction=expr_fraction, with_replacement=with_replacement, shuffle=shuffle, seed=seed
+                    )
                 else:
-                    expr_n = n.expr if hasattr(n, 'expr') else (1 if n is None else n)
-                    res_expr = self.expr.sample(n=expr_n, fraction=None,
-                                                with_replacement=with_replacement,
-                                                shuffle=shuffle, seed=seed)
+                    expr_n = n.expr if hasattr(n, "expr") else (1 if n is None else n)
+                    res_expr = self.expr.sample(
+                        n=expr_n, fraction=None, with_replacement=with_replacement, shuffle=shuffle, seed=seed
+                    )
             except Exception as e:
                 print(f"Warning: Could not create polars expression for list.sample(): {e}")
 
@@ -101,7 +106,7 @@ class ExprListNameSpace:
             fraction=fraction,
             with_replacement=with_replacement,
             shuffle=shuffle,
-            seed=seed
+            seed=seed,
         )
 
     def sum(self) -> Expr:
@@ -134,7 +139,9 @@ class ExprListNameSpace:
 
     def sort(self, *, descending: bool = False, nulls_last: bool = False) -> Expr:
         res_expr = self.expr.sort(descending=descending, nulls_last=nulls_last) if self.expr is not None else None
-        return self._create_next_expr(method_name="sort", result_expr=res_expr, descending=descending, nulls_last=nulls_last)
+        return self._create_next_expr(
+            method_name="sort", result_expr=res_expr, descending=descending, nulls_last=nulls_last
+        )
 
     def reverse(self) -> Expr:
         res_expr = self.expr.reverse() if self.expr is not None else None
@@ -163,7 +170,7 @@ class ExprListNameSpace:
         elif isinstance(other, list):
             if len(other) > 0 and isinstance(other[0], (Expr, str, pl.Series)):
                 # List of expressions
-                other_expr = [o.expr if hasattr(o, 'expr') else (pl.col(o) if isinstance(o, str) else o) for o in other]
+                other_expr = [o.expr if hasattr(o, "expr") else (pl.col(o) if isinstance(o, str) else o) for o in other]
             else:
                 # List of values
                 other_expr = pl.lit(other)
@@ -183,7 +190,7 @@ class ExprListNameSpace:
         return self._create_next_expr(other, method_name="concat", result_expr=res_expr)
 
     def get(self, index: int | Expr | str, *, null_on_oob: bool = False) -> Expr:
-        index_expr = index.expr if hasattr(index, 'expr') else index
+        index_expr = index.expr if hasattr(index, "expr") else index
         res_expr = self.expr.get(index_expr, null_on_oob=null_on_oob) if self.expr is not None else None
         return self._create_next_expr(index, method_name="get", result_expr=res_expr, null_on_oob=null_on_oob)
 
@@ -191,15 +198,15 @@ class ExprListNameSpace:
         indices_expr = indices
         if isinstance(indices, list):
             indices_expr = pl.Series(indices)
-        elif hasattr(indices, 'expr'):
+        elif hasattr(indices, "expr"):
             indices_expr = indices.expr
 
         res_expr = self.expr.gather(indices_expr, null_on_oob=null_on_oob) if self.expr is not None else None
         return self._create_next_expr(indices, method_name="gather", result_expr=res_expr, null_on_oob=null_on_oob)
 
     def gather_every(self, n: int | IntoExprColumn, offset: int | IntoExprColumn = 0) -> Expr:
-        n_expr = n.expr if hasattr(n, 'expr') else n
-        offset_expr = offset.expr if hasattr(offset, 'expr') else offset
+        n_expr = n.expr if hasattr(n, "expr") else n
+        offset_expr = offset.expr if hasattr(offset, "expr") else offset
 
         res_expr = self.expr.gather_every(n_expr, offset_expr) if self.expr is not None else None
         return self._create_next_expr(n, method_name="gather_every", result_expr=res_expr, offset=offset)
@@ -213,12 +220,12 @@ class ExprListNameSpace:
         return self._create_next_expr(method_name="last", result_expr=res_expr)
 
     def contains(self, item: float | str | bool | int | date | datetime | time | IntoExprColumn) -> Expr:
-        item_expr = item.expr if hasattr(item, 'expr') else item
+        item_expr = item.expr if hasattr(item, "expr") else item
         res_expr = self.expr.contains(item_expr) if self.expr is not None else None
         return self._create_next_expr(item, method_name="contains", result_expr=res_expr)
 
     def join(self, separator: IntoExprColumn, *, ignore_nulls: bool = True) -> Expr:
-        separator_expr = separator.expr if hasattr(separator, 'expr') else separator
+        separator_expr = separator.expr if hasattr(separator, "expr") else separator
         res_expr = self.expr.join(separator_expr, ignore_nulls=ignore_nulls) if self.expr is not None else None
         return self._create_next_expr(separator, method_name="join", result_expr=res_expr, ignore_nulls=ignore_nulls)
 
@@ -235,24 +242,24 @@ class ExprListNameSpace:
         return self._create_next_expr(method_name="diff", result_expr=res_expr, n=n, null_behavior=null_behavior)
 
     def shift(self, n: int | IntoExprColumn = 1) -> Expr:
-        n_expr = n.expr if hasattr(n, 'expr') else n
+        n_expr = n.expr if hasattr(n, "expr") else n
         res_expr = self.expr.shift(n_expr) if self.expr is not None else None
         return self._create_next_expr(n, method_name="shift", result_expr=res_expr)
 
     def slice(self, offset: int | str | Expr, length: int | str | Expr | None = None) -> Expr:
-        offset_expr = offset.expr if hasattr(offset, 'expr') else offset
-        length_expr = length.expr if hasattr(length, 'expr') and length is not None else length
+        offset_expr = offset.expr if hasattr(offset, "expr") else offset
+        length_expr = length.expr if hasattr(length, "expr") and length is not None else length
 
         res_expr = self.expr.slice(offset_expr, length_expr) if self.expr is not None else None
         return self._create_next_expr(offset, length, method_name="slice", result_expr=res_expr)
 
     def head(self, n: int | str | Expr = 5) -> Expr:
-        n_expr = n.expr if hasattr(n, 'expr') else n
+        n_expr = n.expr if hasattr(n, "expr") else n
         res_expr = self.expr.head(n_expr) if self.expr is not None else None
         return self._create_next_expr(n, method_name="head", result_expr=res_expr)
 
     def tail(self, n: int | str | Expr = 5) -> Expr:
-        n_expr = n.expr if hasattr(n, 'expr') else n
+        n_expr = n.expr if hasattr(n, "expr") else n
         res_expr = self.expr.tail(n_expr) if self.expr is not None else None
         return self._create_next_expr(n, method_name="tail", result_expr=res_expr)
 
@@ -261,7 +268,7 @@ class ExprListNameSpace:
         return self._create_next_expr(method_name="explode", result_expr=res_expr)
 
     def count_matches(self, element: Any) -> Expr:
-        element_expr = element.expr if hasattr(element, 'expr') else element
+        element_expr = element.expr if hasattr(element, "expr") else element
         res_expr = self.expr.count_matches(element_expr) if self.expr is not None else None
         return self._create_next_expr(element, method_name="count_matches", result_expr=res_expr)
 
@@ -270,10 +277,10 @@ class ExprListNameSpace:
         return self._create_next_expr(width, method_name="to_array", result_expr=res_expr)
 
     def to_struct(
-            self,
-            n_field_strategy: ListToStructWidthStrategy = "first_non_null",
-            fields: Sequence[str] | Callable[[int], str] | None = None,
-            upper_bound: int = 0,
+        self,
+        n_field_strategy: ListToStructWidthStrategy = "first_non_null",
+        fields: Sequence[str] | Callable[[int], str] | None = None,
+        upper_bound: int = 0,
     ) -> Expr:
         res_expr = None
 
@@ -283,9 +290,7 @@ class ExprListNameSpace:
                     res_expr = self.expr.to_struct(fields=fields)
                 else:
                     res_expr = self.expr.to_struct(
-                        n_field_strategy=n_field_strategy,
-                        fields=fields,
-                        upper_bound=upper_bound
+                        n_field_strategy=n_field_strategy, fields=fields, upper_bound=upper_bound
                     )
             except Exception as e:
                 print(f"Warning: Could not create polars expression for list.to_struct(): {e}")
@@ -299,26 +304,26 @@ class ExprListNameSpace:
         )
 
     def eval(self, expr: Expr, *, parallel: bool = False) -> Expr:
-        expr_inner = expr.expr if hasattr(expr, 'expr') else expr
+        expr_inner = expr.expr if hasattr(expr, "expr") else expr
         res_expr = self.expr.eval(expr_inner, parallel=parallel) if self.expr is not None else None
         return self._create_next_expr(expr, method_name="eval", result_expr=res_expr, parallel=parallel)
 
     def set_union(self, other: Any) -> Expr:
-        other_expr = other.expr if hasattr(other, 'expr') else other
+        other_expr = other.expr if hasattr(other, "expr") else other
         res_expr = self.expr.set_union(other_expr) if self.expr is not None else None
         return self._create_next_expr(other, method_name="set_union", result_expr=res_expr)
 
     def set_difference(self, other: Any) -> Expr:
-        other_expr = other.expr if hasattr(other, 'expr') else other
+        other_expr = other.expr if hasattr(other, "expr") else other
         res_expr = self.expr.set_difference(other_expr) if self.expr is not None else None
         return self._create_next_expr(other, method_name="set_difference", result_expr=res_expr)
 
     def set_intersection(self, other: Any) -> Expr:
-        other_expr = other.expr if hasattr(other, 'expr') else other
+        other_expr = other.expr if hasattr(other, "expr") else other
         res_expr = self.expr.set_intersection(other_expr) if self.expr is not None else None
         return self._create_next_expr(other, method_name="set_intersection", result_expr=res_expr)
 
     def set_symmetric_difference(self, other: Any) -> Expr:
-        other_expr = other.expr if hasattr(other, 'expr') else other
+        other_expr = other.expr if hasattr(other, "expr") else other
         res_expr = self.expr.set_symmetric_difference(other_expr) if self.expr is not None else None
-        return self._create_next_expr(other, method_name="set_symmetric_difference", result_expr=res_expr)
+        return self._create_next_expr(other, method_name="set_symmetric_difference", result_expr=res_expr)