Skip to content

Commit 5318254

Browse files
dangotbanneddangotbanned
committed
refactor: Move *Frame 
1 parent 3fc7495 commit 5318254

File tree

5 files changed

+140
-129
lines changed

5 files changed

+140
-129
lines changed

‎narwhals/_plan/arrow/dataframe.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,9 +12,9 @@
1212
from narwhals._plan.arrow import functions as fn
1313
from narwhals._plan.arrow.group_by import ArrowGroupBy as GroupBy
1414
from narwhals._plan.arrow.series import ArrowSeries as Series
15+
from narwhals._plan.compliant.dataframe import EagerDataFrame
1516
from narwhals._plan.compliant.typing import namespace
1617
from narwhals._plan.expressions import NamedIR
17-
from narwhals._plan.protocols import EagerDataFrame
1818
from narwhals._plan.typing import Seq
1919
from narwhals._utils import Version, parse_columns_to_drop
2020
from narwhals.schema import Schema

‎narwhals/_plan/compliant/dataframe.py‎

Lines changed: 131 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1 +1,132 @@
11
from __future__ import annotations
2+
3+
from typing import TYPE_CHECKING, Any, Literal, Protocol, overload
4+
5+
from narwhals._plan.compliant.group_by import Grouped
6+
from narwhals._plan.compliant.typing import ColumnT_co, SeriesT, StoresVersion
7+
from narwhals._plan.typing import (
8+
IntoExpr,
9+
NativeDataFrameT,
10+
NativeFrameT,
11+
NativeSeriesT,
12+
OneOrIterable,
13+
)
14+
15+
if TYPE_CHECKING:
16+
from collections.abc import Iterable, Iterator, Mapping, Sequence
17+
18+
from typing_extensions import Self
19+
20+
from narwhals._plan import expressions as ir
21+
from narwhals._plan.compliant.group_by import (
22+
CompliantGroupBy,
23+
DataFrameGroupBy,
24+
EagerDataFrameGroupBy,
25+
GroupByResolver,
26+
)
27+
from narwhals._plan.compliant.namespace import EagerNamespace
28+
from narwhals._plan.dataframe import BaseFrame, DataFrame
29+
from narwhals._plan.expressions import NamedIR
30+
from narwhals._plan.options import SortMultipleOptions
31+
from narwhals._plan.typing import Seq
32+
from narwhals._utils import Version
33+
from narwhals.dtypes import DType
34+
from narwhals.typing import IntoSchema
35+
36+
37+
class CompliantBaseFrame(StoresVersion, Protocol[ColumnT_co, NativeFrameT]):
38+
_native: NativeFrameT
39+
40+
def __narwhals_namespace__(self) -> Any: ...
41+
@property
42+
def _group_by(self) -> type[CompliantGroupBy[Self]]: ...
43+
@property
44+
def native(self) -> NativeFrameT:
45+
return self._native
46+
47+
@property
48+
def columns(self) -> list[str]: ...
49+
def to_narwhals(self) -> BaseFrame[NativeFrameT]: ...
50+
@classmethod
51+
def from_native(cls, native: NativeFrameT, /, version: Version) -> Self:
52+
obj = cls.__new__(cls)
53+
obj._native = native
54+
obj._version = version
55+
return obj
56+
57+
def _with_native(self, native: NativeFrameT) -> Self:
58+
return self.from_native(native, self.version)
59+
60+
@property
61+
def schema(self) -> Mapping[str, DType]: ...
62+
def _evaluate_irs(
63+
self, nodes: Iterable[NamedIR[ir.ExprIR]], /
64+
) -> Iterator[ColumnT_co]: ...
65+
def select(self, irs: Seq[NamedIR]) -> Self: ...
66+
def select_names(self, *column_names: str) -> Self: ...
67+
def with_columns(self, irs: Seq[NamedIR]) -> Self: ...
68+
def sort(self, by: Seq[NamedIR], options: SortMultipleOptions) -> Self: ...
69+
def drop(self, columns: Sequence[str], *, strict: bool = True) -> Self: ...
70+
def drop_nulls(self, subset: Sequence[str] | None) -> Self: ...
71+
72+
73+
class CompliantDataFrame(
74+
CompliantBaseFrame[SeriesT, NativeDataFrameT],
75+
Protocol[SeriesT, NativeDataFrameT, NativeSeriesT],
76+
):
77+
@property
78+
def _group_by(self) -> type[DataFrameGroupBy[Self]]: ...
79+
@property
80+
def _grouper(self) -> type[Grouped]:
81+
return Grouped
82+
83+
@classmethod
84+
def from_dict(
85+
cls, data: Mapping[str, Any], /, *, schema: IntoSchema | None = None
86+
) -> Self: ...
87+
def group_by_agg(
88+
self, by: OneOrIterable[IntoExpr], aggs: OneOrIterable[IntoExpr], /
89+
) -> Self:
90+
"""Compliant-level `group_by(by).agg(agg)`, allows `Expr`."""
91+
return self._grouper.by(by).agg(aggs).resolve(self).evaluate(self)
92+
93+
def group_by_names(self, names: Seq[str], /) -> DataFrameGroupBy[Self]:
94+
"""Compliant-level `group_by`, allowing only `str` keys."""
95+
return self._group_by.by_names(self, names)
96+
97+
def group_by_resolver(self, resolver: GroupByResolver, /) -> DataFrameGroupBy[Self]:
98+
"""Narwhals-level resolved `group_by`.
99+
100+
`keys`, `aggs` are already parsed and projections planned.
101+
"""
102+
return self._group_by.from_resolver(self, resolver)
103+
104+
def to_narwhals(self) -> DataFrame[NativeDataFrameT, NativeSeriesT]: ...
105+
@overload
106+
def to_dict(self, *, as_series: Literal[True]) -> dict[str, SeriesT]: ...
107+
@overload
108+
def to_dict(self, *, as_series: Literal[False]) -> dict[str, list[Any]]: ...
109+
@overload
110+
def to_dict(
111+
self, *, as_series: bool
112+
) -> dict[str, SeriesT] | dict[str, list[Any]]: ...
113+
def to_dict(
114+
self, *, as_series: bool
115+
) -> dict[str, SeriesT] | dict[str, list[Any]]: ...
116+
def __len__(self) -> int: ...
117+
def with_row_index(self, name: str) -> Self: ...
118+
def row(self, index: int) -> tuple[Any, ...]: ...
119+
120+
121+
class EagerDataFrame(
122+
CompliantDataFrame[SeriesT, NativeDataFrameT, NativeSeriesT],
123+
Protocol[SeriesT, NativeDataFrameT, NativeSeriesT],
124+
):
125+
@property
126+
def _group_by(self) -> type[EagerDataFrameGroupBy[Self]]: ...
127+
def __narwhals_namespace__(self) -> EagerNamespace[Self, SeriesT, Any, Any]: ...
128+
def select(self, irs: Seq[NamedIR]) -> Self:
129+
return self.__narwhals_namespace__()._concat_horizontal(self._evaluate_irs(irs))
130+
131+
def with_columns(self, irs: Seq[NamedIR]) -> Self:
132+
return self.__narwhals_namespace__()._concat_horizontal(self._evaluate_irs(irs))

‎narwhals/_plan/compliant/typing.py‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -7,15 +7,17 @@
77
if TYPE_CHECKING:
88
from typing_extensions import TypeAlias
99

10+
from narwhals._plan.compliant.dataframe import (
11+
CompliantBaseFrame,
12+
CompliantDataFrame,
13+
EagerDataFrame,
14+
)
1015
from narwhals._plan.compliant.group_by import GroupByResolver
1116
from narwhals._plan.compliant.namespace import CompliantNamespace
1217
from narwhals._plan.compliant.series import CompliantSeries
1318
from narwhals._plan.protocols import (
14-
CompliantBaseFrame,
15-
CompliantDataFrame,
1619
CompliantExpr,
1720
CompliantScalar,
18-
EagerDataFrame,
1921
EagerExpr,
2022
EagerScalar,
2123
ExprDispatch,

‎narwhals/_plan/dataframe.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,7 @@
2424
import pyarrow as pa
2525
from typing_extensions import Self
2626

27-
from narwhals._plan.protocols import CompliantBaseFrame, CompliantDataFrame
27+
from narwhals._plan.compliant.dataframe import CompliantBaseFrame, CompliantDataFrame
2828
from narwhals.typing import NativeFrame
2929

3030

‎narwhals/_plan/protocols.py‎

Lines changed: 2 additions & 124 deletions
Original file line numberDiff line numberDiff line change
@@ -2,11 +2,10 @@
22

33
from __future__ import annotations
44

5-
from typing import TYPE_CHECKING, Any, Literal, Protocol, overload
5+
from typing import TYPE_CHECKING, Any, Protocol
66

77
from narwhals._plan.compliant.column import EagerBroadcast, SupportsBroadcast
88
from narwhals._plan.compliant.typing import (
9-
ColumnT_co,
109
FrameT_contra,
1110
LengthT,
1211
NamespaceT_co,
@@ -15,30 +14,12 @@
1514
SeriesT_co,
1615
StoresVersion,
1716
)
18-
from narwhals._plan.typing import (
19-
IntoExpr,
20-
NativeDataFrameT,
21-
NativeFrameT,
22-
NativeSeriesT,
23-
Seq,
24-
)
2517
from narwhals._utils import Version
2618

2719
if TYPE_CHECKING:
28-
from collections.abc import Iterable, Iterator, Mapping, Sequence
29-
3020
from typing_extensions import Self
3121

3222
from narwhals._plan import expressions as ir
33-
from narwhals._plan.compliant.group_by import (
34-
CompliantGroupBy,
35-
DataFrameGroupBy,
36-
EagerDataFrameGroupBy,
37-
GroupByResolver,
38-
Grouped,
39-
)
40-
from narwhals._plan.compliant.namespace import EagerNamespace
41-
from narwhals._plan.dataframe import BaseFrame, DataFrame
4223
from narwhals._plan.expressions import (
4324
BinaryExpr,
4425
FunctionExpr,
@@ -48,10 +29,7 @@
4829
functions as F,
4930
)
5031
from narwhals._plan.expressions.boolean import IsBetween, IsFinite, IsNan, IsNull, Not
51-
from narwhals._plan.options import SortMultipleOptions
52-
from narwhals._plan.typing import OneOrIterable
53-
from narwhals.dtypes import DType
54-
from narwhals.typing import IntoDType, IntoSchema, PythonLiteral
32+
from narwhals.typing import IntoDType, PythonLiteral
5533

5634

5735
class ExprDispatch(StoresVersion, Protocol[FrameT_contra, R_co, NamespaceT_co]):
@@ -307,103 +285,3 @@ class LazyScalar(
307285
LazyExpr[FrameT_contra, SeriesT, LengthT],
308286
Protocol[FrameT_contra, SeriesT, LengthT],
309287
): ...
310-
311-
312-
class CompliantBaseFrame(StoresVersion, Protocol[ColumnT_co, NativeFrameT]):
313-
_native: NativeFrameT
314-
315-
def __narwhals_namespace__(self) -> Any: ...
316-
@property
317-
def _group_by(self) -> type[CompliantGroupBy[Self]]: ...
318-
@property
319-
def native(self) -> NativeFrameT:
320-
return self._native
321-
322-
@property
323-
def columns(self) -> list[str]: ...
324-
def to_narwhals(self) -> BaseFrame[NativeFrameT]: ...
325-
@classmethod
326-
def from_native(cls, native: NativeFrameT, /, version: Version) -> Self:
327-
obj = cls.__new__(cls)
328-
obj._native = native
329-
obj._version = version
330-
return obj
331-
332-
def _with_native(self, native: NativeFrameT) -> Self:
333-
return self.from_native(native, self.version)
334-
335-
@property
336-
def schema(self) -> Mapping[str, DType]: ...
337-
def _evaluate_irs(
338-
self, nodes: Iterable[NamedIR[ir.ExprIR]], /
339-
) -> Iterator[ColumnT_co]: ...
340-
def select(self, irs: Seq[NamedIR]) -> Self: ...
341-
def select_names(self, *column_names: str) -> Self: ...
342-
def with_columns(self, irs: Seq[NamedIR]) -> Self: ...
343-
def sort(self, by: Seq[NamedIR], options: SortMultipleOptions) -> Self: ...
344-
def drop(self, columns: Sequence[str], *, strict: bool = True) -> Self: ...
345-
def drop_nulls(self, subset: Sequence[str] | None) -> Self: ...
346-
347-
348-
class CompliantDataFrame(
349-
CompliantBaseFrame[SeriesT, NativeDataFrameT],
350-
Protocol[SeriesT, NativeDataFrameT, NativeSeriesT],
351-
):
352-
@property
353-
def _group_by(self) -> type[DataFrameGroupBy[Self]]: ...
354-
@property
355-
def _grouper(self) -> type[Grouped]:
356-
from narwhals._plan.compliant.group_by import Grouped
357-
358-
return Grouped
359-
360-
@classmethod
361-
def from_dict(
362-
cls, data: Mapping[str, Any], /, *, schema: IntoSchema | None = None
363-
) -> Self: ...
364-
def group_by_agg(
365-
self, by: OneOrIterable[IntoExpr], aggs: OneOrIterable[IntoExpr], /
366-
) -> Self:
367-
"""Compliant-level `group_by(by).agg(agg)`, allows `Expr`."""
368-
return self._grouper.by(by).agg(aggs).resolve(self).evaluate(self)
369-
370-
def group_by_names(self, names: Seq[str], /) -> DataFrameGroupBy[Self]:
371-
"""Compliant-level `group_by`, allowing only `str` keys."""
372-
return self._group_by.by_names(self, names)
373-
374-
def group_by_resolver(self, resolver: GroupByResolver, /) -> DataFrameGroupBy[Self]:
375-
"""Narwhals-level resolved `group_by`.
376-
377-
`keys`, `aggs` are already parsed and projections planned.
378-
"""
379-
return self._group_by.from_resolver(self, resolver)
380-
381-
def to_narwhals(self) -> DataFrame[NativeDataFrameT, NativeSeriesT]: ...
382-
@overload
383-
def to_dict(self, *, as_series: Literal[True]) -> dict[str, SeriesT]: ...
384-
@overload
385-
def to_dict(self, *, as_series: Literal[False]) -> dict[str, list[Any]]: ...
386-
@overload
387-
def to_dict(
388-
self, *, as_series: bool
389-
) -> dict[str, SeriesT] | dict[str, list[Any]]: ...
390-
def to_dict(
391-
self, *, as_series: bool
392-
) -> dict[str, SeriesT] | dict[str, list[Any]]: ...
393-
def __len__(self) -> int: ...
394-
def with_row_index(self, name: str) -> Self: ...
395-
def row(self, index: int) -> tuple[Any, ...]: ...
396-
397-
398-
class EagerDataFrame(
399-
CompliantDataFrame[SeriesT, NativeDataFrameT, NativeSeriesT],
400-
Protocol[SeriesT, NativeDataFrameT, NativeSeriesT],
401-
):
402-
@property
403-
def _group_by(self) -> type[EagerDataFrameGroupBy[Self]]: ...
404-
def __narwhals_namespace__(self) -> EagerNamespace[Self, SeriesT, Any, Any]: ...
405-
def select(self, irs: Seq[NamedIR]) -> Self:
406-
return self.__narwhals_namespace__()._concat_horizontal(self._evaluate_irs(irs))
407-
408-
def with_columns(self, irs: Seq[NamedIR]) -> Self:
409-
return self.__narwhals_namespace__()._concat_horizontal(self._evaluate_irs(irs))

0 commit comments

Comments
 (0)