Skip to content

Commit a11807a

Browse files
committed
tie IntoFrameT to NativeFrame
1 parent 45d0231 commit a11807a

File tree

4 files changed

+35
-37
lines changed

4 files changed

+35
-37
lines changed

narwhals/_compliant/dataframe.py

Lines changed: 11 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -12,8 +12,7 @@
1212
CompliantSeriesT,
1313
EagerExprT,
1414
EagerSeriesT,
15-
NativeDataFrameT,
16-
NativeLazyFrameT,
15+
NativeFrameT,
1716
NativeSeriesT,
1817
)
1918
from narwhals._translate import (
@@ -179,8 +178,8 @@ class CompliantDataFrame(
179178
DictConvertible["_ToDict[CompliantSeriesT]", Mapping[str, Any]],
180179
ArrowConvertible["pa.Table", "IntoArrowTable"],
181180
Sized,
182-
CompliantFrame[CompliantExprT_contra, NativeDataFrameT, ToNarwhalsT_co],
183-
Protocol[CompliantSeriesT, CompliantExprT_contra, NativeDataFrameT, ToNarwhalsT_co],
181+
CompliantFrame[CompliantExprT_contra, NativeFrameT, ToNarwhalsT_co],
182+
Protocol[CompliantSeriesT, CompliantExprT_contra, NativeFrameT, ToNarwhalsT_co],
184183
):
185184
def __narwhals_dataframe__(self) -> Self: ...
186185
@classmethod
@@ -292,8 +291,8 @@ def write_parquet(self, file: str | Path | BytesIO) -> None: ...
292291

293292

294293
class CompliantLazyFrame(
295-
CompliantFrame[CompliantExprT_contra, NativeLazyFrameT, ToNarwhalsT_co],
296-
Protocol[CompliantExprT_contra, NativeLazyFrameT, ToNarwhalsT_co],
294+
CompliantFrame[CompliantExprT_contra, NativeFrameT, ToNarwhalsT_co],
295+
Protocol[CompliantExprT_contra, NativeFrameT, ToNarwhalsT_co],
297296
):
298297
def __narwhals_lazyframe__(self) -> Self: ...
299298
# `LazySelectorNamespace._iter_columns` depends
@@ -312,24 +311,20 @@ def sink_parquet(self, file: str | Path | BytesIO) -> None: ...
312311

313312

314313
class EagerDataFrame(
315-
CompliantDataFrame[
316-
EagerSeriesT, EagerExprT, NativeDataFrameT, "DataFrame[NativeDataFrameT]"
317-
],
318-
CompliantLazyFrame[EagerExprT, "Incomplete", "DataFrame[NativeDataFrameT]"],
314+
CompliantDataFrame[EagerSeriesT, EagerExprT, NativeFrameT, "DataFrame[NativeFrameT]"],
315+
CompliantLazyFrame[EagerExprT, "Incomplete", "DataFrame[NativeFrameT]"],
319316
ValidateBackendVersion,
320-
Protocol[EagerSeriesT, EagerExprT, NativeDataFrameT, NativeSeriesT],
317+
Protocol[EagerSeriesT, EagerExprT, NativeFrameT, NativeSeriesT],
321318
):
322319
@property
323320
def _backend_version(self) -> tuple[int, ...]:
324321
return self._implementation._backend_version()
325322

326323
def __narwhals_namespace__(
327324
self,
328-
) -> EagerNamespace[
329-
Self, EagerSeriesT, EagerExprT, NativeDataFrameT, NativeSeriesT
330-
]: ...
325+
) -> EagerNamespace[Self, EagerSeriesT, EagerExprT, NativeFrameT, NativeSeriesT]: ...
331326

332-
def to_narwhals(self) -> DataFrame[NativeDataFrameT]:
327+
def to_narwhals(self) -> DataFrame[NativeFrameT]:
333328
return self._version.dataframe(self, level="full")
334329

335330
def aggregate(self, *exprs: EagerExprT) -> Self: # pyright: ignore[reportIncompatibleMethodOverride]
@@ -343,7 +338,7 @@ def aggregate(self, *exprs: EagerExprT) -> Self: # pyright: ignore[reportIncomp
343338
return self.select(*exprs) # pyright: ignore[reportArgumentType]
344339

345340
def _with_native(
346-
self, df: NativeDataFrameT, *, validate_column_names: bool = True
341+
self, df: NativeFrameT, *, validate_column_names: bool = True
347342
) -> Self: ...
348343

349344
def _check_columns_exist(self, subset: Sequence[str]) -> ColumnNotFoundError | None:

narwhals/_native.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -287,7 +287,7 @@ def dropDuplicatesWithinWatermark(self, *arg: Any, **kwargs: Any) -> Any: ... #
287287
NativePySparkConnect: TypeAlias = _PySparkDataFrame
288288
NativeSparkLike: TypeAlias = "NativeSQLFrame | NativePySpark | NativePySparkConnect"
289289
NativeKnown: TypeAlias = "NativePolars | NativeArrow | NativePandasLike | NativeSparkLike | NativeDuckDB | NativeDask | NativeIbis"
290-
NativeUnknown: TypeAlias = "NativeDataFrame | NativeSeries | NativeLazyFrame"
290+
NativeUnknown: TypeAlias = "NativeFrame | NativeSeries"
291291
NativeAny: TypeAlias = "NativeKnown | NativeUnknown"
292292

293293
IntoDataFrame: TypeAlias = NativeDataFrame
@@ -304,7 +304,7 @@ def dropDuplicatesWithinWatermark(self, *arg: Any, **kwargs: Any) -> Any: ... #
304304
"""
305305

306306
IntoLazyFrame: TypeAlias = Union[NativeLazyFrame, NativeIbis]
307-
IntoFrame: TypeAlias = Union[IntoDataFrame, IntoLazyFrame]
307+
IntoFrame: TypeAlias = NativeFrame
308308
"""Anything which can be converted to a Narwhals DataFrame or LazyFrame.
309309
310310
Use this if your function can accept an object which can be converted to either

narwhals/dataframe.py

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -95,7 +95,7 @@
9595
PS = ParamSpec("PS")
9696
Incomplete: TypeAlias = Any
9797

98-
_FrameT = TypeVar("_FrameT", bound="IntoFrame")
98+
FrameT = TypeVar("FrameT", bound="IntoFrame")
9999
LazyFrameT = TypeVar("LazyFrameT", bound="IntoLazyFrame")
100100
DataFrameT = TypeVar("DataFrameT", bound="IntoDataFrame")
101101
R = TypeVar("R")
@@ -104,7 +104,7 @@
104104
MultiIndexSelector: TypeAlias = "_MultiIndexSelector[Series[Any]]"
105105

106106

107-
class BaseFrame(Generic[_FrameT]):
107+
class BaseFrame(Generic[FrameT]):
108108
_compliant_frame: Any
109109
_level: Literal["full", "lazy", "interchange"]
110110

@@ -447,7 +447,7 @@ def explode(self, columns: str | Sequence[str], *more_columns: str) -> Self:
447447
return self._with_compliant(self._compliant_frame.explode(columns=to_explode))
448448

449449

450-
class DataFrame(BaseFrame[DataFrameT]):
450+
class DataFrame(BaseFrame[FrameT]):
451451
"""Narwhals DataFrame, backed by a native eager dataframe.
452452
453453
Warning:
@@ -475,7 +475,7 @@ class DataFrame(BaseFrame[DataFrameT]):
475475
_version: ClassVar[Version] = Version.MAIN
476476

477477
@property
478-
def _compliant(self) -> CompliantDataFrame[Any, Any, DataFrameT, Self]:
478+
def _compliant(self) -> CompliantDataFrame[Any, Any, FrameT, Self]:
479479
return self._compliant_frame
480480

481481
@property
@@ -492,7 +492,7 @@ def _validate_metadata(self, metadata: ExprMetadata) -> None:
492492

493493
def __init__(self, df: Any, *, level: Literal["full", "lazy", "interchange"]) -> None:
494494
self._level: Literal["full", "lazy", "interchange"] = level
495-
self._compliant_frame: CompliantDataFrame[Any, Any, DataFrameT, Self]
495+
self._compliant_frame: CompliantDataFrame[Any, Any, FrameT, Self]
496496
if is_compliant_dataframe(df):
497497
self._compliant_frame = df.__narwhals_dataframe__()
498498
else: # pragma: no cover
@@ -871,7 +871,7 @@ def lazy(
871871
msg = f"Not-supported backend.\n\nExpected one of {get_args(_LazyAllowedImpl)} or `None`, got {lazy_backend}"
872872
raise ValueError(msg)
873873

874-
def to_native(self) -> DataFrameT:
874+
def to_native(self) -> FrameT:
875875
"""Convert Narwhals DataFrame to native one.
876876
877877
Examples:
@@ -2348,7 +2348,7 @@ def explode(self, columns: str | Sequence[str], *more_columns: str) -> Self:
23482348
return super().explode(columns, *more_columns)
23492349

23502350

2351-
class LazyFrame(BaseFrame[LazyFrameT]):
2351+
class LazyFrame(BaseFrame[FrameT]):
23522352
"""Narwhals LazyFrame, backed by a native lazyframe.
23532353
23542354
Warning:
@@ -2362,7 +2362,7 @@ class LazyFrame(BaseFrame[LazyFrameT]):
23622362
"""
23632363

23642364
@property
2365-
def _compliant(self) -> CompliantLazyFrame[Any, LazyFrameT, Self]:
2365+
def _compliant(self) -> CompliantLazyFrame[Any, FrameT, Self]:
23662366
return self._compliant_frame
23672367

23682368
@property
@@ -2395,7 +2395,7 @@ def _validate_metadata(self, metadata: ExprMetadata) -> None:
23952395

23962396
def __init__(self, df: Any, *, level: Literal["full", "lazy", "interchange"]) -> None:
23972397
self._level = level
2398-
self._compliant_frame: CompliantLazyFrame[Any, LazyFrameT, Self]
2398+
self._compliant_frame: CompliantLazyFrame[Any, FrameT, Self]
23992399
if is_compliant_lazyframe(df):
24002400
self._compliant_frame = df.__narwhals_lazyframe__()
24012401
else: # pragma: no cover
@@ -2478,7 +2478,7 @@ def collect(
24782478
msg = f"Unsupported `backend` value.\nExpected one of {get_args(_LazyFrameCollectImpl)} or None, got: {eager_backend}."
24792479
raise ValueError(msg)
24802480

2481-
def to_native(self) -> LazyFrameT:
2481+
def to_native(self) -> FrameT:
24822482
"""Convert Narwhals LazyFrame to native one.
24832483
24842484
Examples:

narwhals/translate.py

Lines changed: 12 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,7 @@
5555
Frame,
5656
IntoDataFrameT,
5757
IntoFrame,
58+
IntoFrameT,
5859
IntoLazyFrameT,
5960
IntoSeries,
6061
IntoSeriesT,
@@ -70,12 +71,12 @@
7071

7172
@overload
7273
def to_native(
73-
narwhals_object: DataFrame[IntoDataFrameT], *, pass_through: Literal[False] = ...
74-
) -> IntoDataFrameT: ...
74+
narwhals_object: DataFrame[IntoFrameT], *, pass_through: Literal[False] = ...
75+
) -> IntoFrameT: ...
7576
@overload
7677
def to_native(
77-
narwhals_object: LazyFrame[IntoLazyFrameT], *, pass_through: Literal[False] = ...
78-
) -> IntoLazyFrameT: ...
78+
narwhals_object: LazyFrame[IntoFrameT], *, pass_through: Literal[False] = ...
79+
) -> IntoFrameT: ...
7980
@overload
8081
def to_native(
8182
narwhals_object: Series[IntoSeriesT], *, pass_through: Literal[False] = ...
@@ -85,12 +86,10 @@ def to_native(narwhals_object: Any, *, pass_through: bool) -> Any: ...
8586

8687

8788
def to_native(
88-
narwhals_object: DataFrame[IntoDataFrameT]
89-
| LazyFrame[IntoLazyFrameT]
90-
| Series[IntoSeriesT],
89+
narwhals_object: DataFrame[IntoFrameT] | LazyFrame[IntoFrameT] | Series[IntoSeriesT],
9190
*,
9291
pass_through: bool = False,
93-
) -> IntoDataFrameT | IntoLazyFrameT | IntoSeriesT | Any:
92+
) -> IntoFrameT | IntoSeriesT | Any:
9493
"""Convert Narwhals object to native one.
9594
9695
Arguments:
@@ -145,6 +144,10 @@ def from_native(
145144
native_object: IntoLazyFrameT, **kwds: Unpack[AllowLazy]
146145
) -> LazyFrame[IntoLazyFrameT]: ...
147146
@overload
147+
def from_native( # type: ignore[overload-overlap]
148+
native_object: IntoFrameT, **kwds: Unpack[AllowLazy]
149+
) -> DataFrame[IntoFrameT] | LazyFrame[IntoFrameT]: ...
150+
@overload
148151
def from_native(
149152
native_object: IntoDataFrameT | IntoSeriesT, **kwds: Unpack[AllowSeries]
150153
) -> DataFrame[IntoDataFrameT] | Series[IntoSeriesT]: ...
@@ -164,7 +167,7 @@ def from_native(
164167
series_only: bool,
165168
allow_series: bool | None,
166169
) -> Any: ...
167-
def from_native(
170+
def from_native( # pyright: ignore[reportInconsistentOverload]
168171
native_object: IntoLazyFrameT
169172
| IntoDataFrameT
170173
| IntoSeriesT

0 commit comments

Comments
 (0)