|
12 | 12 | from typing import Literal |
13 | 13 | from typing import Sequence |
14 | 14 | from typing import TypeVar |
15 | | -from typing import overload |
16 | 15 |
|
17 | 16 | from narwhals.dependencies import is_narwhals_series |
18 | 17 | from narwhals.dependencies import is_numpy_array |
|
30 | 29 | from narwhals._compliant import CompliantNamespace |
31 | 30 | from narwhals._compliant import CompliantSeriesOrNativeExprT_co |
32 | 31 | from narwhals._compliant import CompliantSeriesT_co |
33 | | - from narwhals._pandas_like.expr import PandasLikeExpr |
34 | 32 | from narwhals.expr import Expr |
35 | 33 | from narwhals.typing import CompliantDataFrame |
36 | 34 | from narwhals.typing import CompliantLazyFrame |
37 | | - from narwhals.typing import CompliantSeries |
38 | 35 | from narwhals.typing import IntoExpr |
39 | 36 | from narwhals.typing import _1DArray |
40 | 37 |
|
41 | | - PandasLikeExprT = TypeVar("PandasLikeExprT", bound=PandasLikeExpr) |
42 | | - |
43 | 38 | T = TypeVar("T") |
44 | 39 |
|
45 | 40 |
|
@@ -79,122 +74,6 @@ def evaluate_into_exprs( |
79 | 74 | return list(chain.from_iterable(evaluate_into_expr(df, expr) for expr in exprs)) |
80 | 75 |
|
81 | 76 |
|
82 | | -@overload |
83 | | -def maybe_evaluate_expr( |
84 | | - df: CompliantFrameT, expr: CompliantExpr[CompliantFrameT, CompliantSeriesT_co] |
85 | | -) -> CompliantSeriesT_co: ... |
86 | | - |
87 | | - |
88 | | -@overload |
89 | | -def maybe_evaluate_expr(df: CompliantDataFrame[Any], expr: T) -> T: ... |
90 | | - |
91 | | - |
92 | | -def maybe_evaluate_expr( |
93 | | - df: Any, expr: CompliantExpr[Any, CompliantSeriesT_co] | T |
94 | | -) -> CompliantSeriesT_co | T: |
95 | | - """Evaluate `expr` if it's an expression, otherwise return it as is.""" |
96 | | - if is_compliant_expr(expr): |
97 | | - result: Sequence[CompliantSeriesT_co] = expr(df) |
98 | | - if len(result) > 1: |
99 | | - msg = "Multi-output expressions (e.g. `nw.all()` or `nw.col('a', 'b')`) are not supported in this context" |
100 | | - raise ValueError(msg) |
101 | | - return result[0] |
102 | | - return expr |
103 | | - |
104 | | - |
105 | | -def reuse_series_implementation( |
106 | | - expr: PandasLikeExprT, |
107 | | - attr: str, |
108 | | - *, |
109 | | - returns_scalar: bool = False, |
110 | | - call_kwargs: dict[str, Any] | None = None, |
111 | | - **expressifiable_args: Any, |
112 | | -) -> PandasLikeExprT: |
113 | | - """Reuse Series implementation for expression. |
114 | | -
|
115 | | - If Series.foo is already defined, and we'd like Expr.foo to be the same, we can |
116 | | - leverage this method to do that for us. |
117 | | -
|
118 | | - Arguments: |
119 | | - expr: expression object. |
120 | | - attr: name of method. |
121 | | - returns_scalar: whether the Series version returns a scalar. In this case, |
122 | | - the expression version should return a 1-row Series. |
123 | | - call_kwargs: non-expressifiable args which we may need to reuse in `agg` or `over`, |
124 | | - such as `ddof` for `std` and `var`. |
125 | | - expressifiable_args: keyword arguments to pass to function, which may |
126 | | - be expressifiable (e.g. `nw.col('a').is_between(3, nw.col('b')))`). |
127 | | - """ |
128 | | - plx = expr.__narwhals_namespace__() |
129 | | - |
130 | | - def func(df: CompliantDataFrame[Any]) -> Sequence[CompliantSeries]: |
131 | | - _kwargs = { |
132 | | - **(call_kwargs or {}), |
133 | | - **{ |
134 | | - arg_name: maybe_evaluate_expr(df, arg_value) |
135 | | - for arg_name, arg_value in expressifiable_args.items() |
136 | | - }, |
137 | | - } |
138 | | - |
139 | | - out: list[CompliantSeries] = [ |
140 | | - plx._create_series_from_scalar( # type: ignore # noqa: PGH003 |
141 | | - getattr(series, attr)(**_kwargs), |
142 | | - reference_series=series, |
143 | | - ) |
144 | | - if returns_scalar |
145 | | - else getattr(series, attr)(**_kwargs) |
146 | | - for series in expr(df) # type: ignore # noqa: PGH003 |
147 | | - ] |
148 | | - _, aliases = evaluate_output_names_and_aliases(expr, df, []) |
149 | | - if [s.name for s in out] != list(aliases): # pragma: no cover |
150 | | - msg = ( |
151 | | - f"Safety assertion failed, please report a bug to https://github.com/narwhals-dev/narwhals/issues\n" |
152 | | - f"Expression aliases: {aliases}\n" |
153 | | - f"Series names: {[s.name for s in out]}" |
154 | | - ) |
155 | | - raise AssertionError(msg) |
156 | | - return out |
157 | | - |
158 | | - return plx._create_expr_from_callable( # type: ignore # noqa: PGH003 |
159 | | - func, |
160 | | - depth=expr._depth + 1, |
161 | | - function_name=f"{expr._function_name}->{attr}", |
162 | | - evaluate_output_names=expr._evaluate_output_names, |
163 | | - alias_output_names=expr._alias_output_names, |
164 | | - call_kwargs=call_kwargs, |
165 | | - ) |
166 | | - |
167 | | - |
168 | | -def reuse_series_namespace_implementation( |
169 | | - expr: PandasLikeExprT, |
170 | | - series_namespace: str, |
171 | | - attr: str, |
172 | | - **kwargs: Any, |
173 | | -) -> PandasLikeExprT: |
174 | | - """Reuse Series implementation for expression. |
175 | | -
|
176 | | - Just like `reuse_series_implementation`, but for e.g. `Expr.dt.foo` instead |
177 | | - of `Expr.foo`. |
178 | | -
|
179 | | - Arguments: |
180 | | - expr: expression object. |
181 | | - series_namespace: The Series namespace (e.g. `dt`, `cat`, `str`, `list`, `name`) |
182 | | - attr: name of method. |
183 | | - kwargs: keyword arguments to pass to function. |
184 | | - """ |
185 | | - plx = expr.__narwhals_namespace__() |
186 | | - return plx._create_expr_from_callable( # type: ignore # noqa: PGH003 |
187 | | - lambda df: [ |
188 | | - getattr(getattr(series, series_namespace), attr)(**kwargs) |
189 | | - for series in expr(df) |
190 | | - ], |
191 | | - depth=expr._depth + 1, |
192 | | - function_name=f"{expr._function_name}->{series_namespace}.{attr}", |
193 | | - evaluate_output_names=expr._evaluate_output_names, |
194 | | - alias_output_names=expr._alias_output_names, |
195 | | - ) |
196 | | - |
197 | | - |
198 | 77 | def is_elementary_expression(expr: CompliantExpr[Any, Any]) -> bool: |
199 | 78 | """Check if expr is elementary. |
200 | 79 |
|
|
0 commit comments