diff --git a/narwhals/_compliant/dataframe.py b/narwhals/_compliant/dataframe.py index 3e7810616c..3249947553 100644 --- a/narwhals/_compliant/dataframe.py +++ b/narwhals/_compliant/dataframe.py @@ -12,8 +12,7 @@ CompliantSeriesT, EagerExprT, EagerSeriesT, - NativeDataFrameT, - NativeLazyFrameT, + NativeFrameT, NativeSeriesT, ) from narwhals._translate import ( @@ -179,8 +178,8 @@ class CompliantDataFrame( DictConvertible["_ToDict[CompliantSeriesT]", Mapping[str, Any]], ArrowConvertible["pa.Table", "IntoArrowTable"], Sized, - CompliantFrame[CompliantExprT_contra, NativeDataFrameT, ToNarwhalsT_co], - Protocol[CompliantSeriesT, CompliantExprT_contra, NativeDataFrameT, ToNarwhalsT_co], + CompliantFrame[CompliantExprT_contra, NativeFrameT, ToNarwhalsT_co], + Protocol[CompliantSeriesT, CompliantExprT_contra, NativeFrameT, ToNarwhalsT_co], ): def __narwhals_dataframe__(self) -> Self: ... @classmethod @@ -292,8 +291,8 @@ def write_parquet(self, file: str | Path | BytesIO) -> None: ... class CompliantLazyFrame( - CompliantFrame[CompliantExprT_contra, NativeLazyFrameT, ToNarwhalsT_co], - Protocol[CompliantExprT_contra, NativeLazyFrameT, ToNarwhalsT_co], + CompliantFrame[CompliantExprT_contra, NativeFrameT, ToNarwhalsT_co], + Protocol[CompliantExprT_contra, NativeFrameT, ToNarwhalsT_co], ): def __narwhals_lazyframe__(self) -> Self: ... # `LazySelectorNamespace._iter_columns` depends @@ -312,12 +311,10 @@ def sink_parquet(self, file: str | Path | BytesIO) -> None: ... class EagerDataFrame( - CompliantDataFrame[ - EagerSeriesT, EagerExprT, NativeDataFrameT, "DataFrame[NativeDataFrameT]" - ], - CompliantLazyFrame[EagerExprT, "Incomplete", "DataFrame[NativeDataFrameT]"], + CompliantDataFrame[EagerSeriesT, EagerExprT, NativeFrameT, "DataFrame[NativeFrameT]"], + CompliantLazyFrame[EagerExprT, "Incomplete", "DataFrame[NativeFrameT]"], ValidateBackendVersion, - Protocol[EagerSeriesT, EagerExprT, NativeDataFrameT, NativeSeriesT], + Protocol[EagerSeriesT, EagerExprT, NativeFrameT, NativeSeriesT], ): @property def _backend_version(self) -> tuple[int, ...]: @@ -325,11 +322,9 @@ def _backend_version(self) -> tuple[int, ...]: def __narwhals_namespace__( self, - ) -> EagerNamespace[ - Self, EagerSeriesT, EagerExprT, NativeDataFrameT, NativeSeriesT - ]: ... + ) -> EagerNamespace[Self, EagerSeriesT, EagerExprT, NativeFrameT, NativeSeriesT]: ... - def to_narwhals(self) -> DataFrame[NativeDataFrameT]: + def to_narwhals(self) -> DataFrame[NativeFrameT]: return self._version.dataframe(self, level="full") def aggregate(self, *exprs: EagerExprT) -> Self: # pyright: ignore[reportIncompatibleMethodOverride] @@ -343,7 +338,7 @@ def aggregate(self, *exprs: EagerExprT) -> Self: # pyright: ignore[reportIncomp return self.select(*exprs) # pyright: ignore[reportArgumentType] def _with_native( - self, df: NativeDataFrameT, *, validate_column_names: bool = True + self, df: NativeFrameT, *, validate_column_names: bool = True ) -> Self: ... def _check_columns_exist(self, subset: Sequence[str]) -> ColumnNotFoundError | None: diff --git a/narwhals/_native.py b/narwhals/_native.py index 2e06ac7041..96d5c45d03 100644 --- a/narwhals/_native.py +++ b/narwhals/_native.py @@ -287,7 +287,7 @@ def dropDuplicatesWithinWatermark(self, *arg: Any, **kwargs: Any) -> Any: ... # NativePySparkConnect: TypeAlias = _PySparkDataFrame NativeSparkLike: TypeAlias = "NativeSQLFrame | NativePySpark | NativePySparkConnect" NativeKnown: TypeAlias = "NativePolars | NativeArrow | NativePandasLike | NativeSparkLike | NativeDuckDB | NativeDask | NativeIbis" -NativeUnknown: TypeAlias = "NativeDataFrame | NativeSeries | NativeLazyFrame" +NativeUnknown: TypeAlias = "NativeFrame | NativeSeries" NativeAny: TypeAlias = "NativeKnown | NativeUnknown" IntoDataFrame: TypeAlias = NativeDataFrame @@ -304,7 +304,7 @@ def dropDuplicatesWithinWatermark(self, *arg: Any, **kwargs: Any) -> Any: ... # """ IntoLazyFrame: TypeAlias = Union[NativeLazyFrame, NativeIbis] -IntoFrame: TypeAlias = Union[IntoDataFrame, IntoLazyFrame] +IntoFrame: TypeAlias = NativeFrame """Anything which can be converted to a Narwhals DataFrame or LazyFrame. Use this if your function can accept an object which can be converted to either diff --git a/narwhals/dataframe.py b/narwhals/dataframe.py index bd602efd52..2268c10142 100644 --- a/narwhals/dataframe.py +++ b/narwhals/dataframe.py @@ -95,7 +95,7 @@ PS = ParamSpec("PS") Incomplete: TypeAlias = Any -_FrameT = TypeVar("_FrameT", bound="IntoFrame") +FrameT = TypeVar("FrameT", bound="IntoFrame") LazyFrameT = TypeVar("LazyFrameT", bound="IntoLazyFrame") DataFrameT = TypeVar("DataFrameT", bound="IntoDataFrame") R = TypeVar("R") @@ -104,7 +104,7 @@ MultiIndexSelector: TypeAlias = "_MultiIndexSelector[Series[Any]]" -class BaseFrame(Generic[_FrameT]): +class BaseFrame(Generic[FrameT]): _compliant_frame: Any _level: Literal["full", "lazy", "interchange"] @@ -447,7 +447,7 @@ def explode(self, columns: str | Sequence[str], *more_columns: str) -> Self: return self._with_compliant(self._compliant_frame.explode(columns=to_explode)) -class DataFrame(BaseFrame[DataFrameT]): +class DataFrame(BaseFrame[FrameT]): """Narwhals DataFrame, backed by a native eager dataframe. Warning: @@ -475,7 +475,7 @@ class DataFrame(BaseFrame[DataFrameT]): _version: ClassVar[Version] = Version.MAIN @property - def _compliant(self) -> CompliantDataFrame[Any, Any, DataFrameT, Self]: + def _compliant(self) -> CompliantDataFrame[Any, Any, FrameT, Self]: return self._compliant_frame @property @@ -492,7 +492,7 @@ def _validate_metadata(self, metadata: ExprMetadata) -> None: def __init__(self, df: Any, *, level: Literal["full", "lazy", "interchange"]) -> None: self._level: Literal["full", "lazy", "interchange"] = level - self._compliant_frame: CompliantDataFrame[Any, Any, DataFrameT, Self] + self._compliant_frame: CompliantDataFrame[Any, Any, FrameT, Self] if is_compliant_dataframe(df): self._compliant_frame = df.__narwhals_dataframe__() else: # pragma: no cover @@ -871,7 +871,7 @@ def lazy( msg = f"Not-supported backend.\n\nExpected one of {get_args(_LazyAllowedImpl)} or `None`, got {lazy_backend}" raise ValueError(msg) - def to_native(self) -> DataFrameT: + def to_native(self) -> FrameT: """Convert Narwhals DataFrame to native one. Examples: @@ -2348,7 +2348,7 @@ def explode(self, columns: str | Sequence[str], *more_columns: str) -> Self: return super().explode(columns, *more_columns) -class LazyFrame(BaseFrame[LazyFrameT]): +class LazyFrame(BaseFrame[FrameT]): """Narwhals LazyFrame, backed by a native lazyframe. Warning: @@ -2364,7 +2364,7 @@ class LazyFrame(BaseFrame[LazyFrameT]): _version: ClassVar[Version] = Version.MAIN @property - def _compliant(self) -> CompliantLazyFrame[Any, LazyFrameT, Self]: + def _compliant(self) -> CompliantLazyFrame[Any, FrameT, Self]: return self._compliant_frame @property @@ -2397,7 +2397,7 @@ def _validate_metadata(self, metadata: ExprMetadata) -> None: def __init__(self, df: Any, *, level: Literal["full", "lazy", "interchange"]) -> None: self._level = level - self._compliant_frame: CompliantLazyFrame[Any, LazyFrameT, Self] + self._compliant_frame: CompliantLazyFrame[Any, FrameT, Self] if is_compliant_lazyframe(df): self._compliant_frame = df.__narwhals_lazyframe__() else: # pragma: no cover @@ -2480,7 +2480,7 @@ def collect( msg = f"Unsupported `backend` value.\nExpected one of {get_args(_LazyFrameCollectImpl)} or None, got: {eager_backend}." raise ValueError(msg) - def to_native(self) -> LazyFrameT: + def to_native(self) -> FrameT: """Convert Narwhals LazyFrame to native one. Examples: diff --git a/narwhals/stable/v2/__init__.py b/narwhals/stable/v2/__init__.py index dc4c757c20..285147f596 100644 --- a/narwhals/stable/v2/__init__.py +++ b/narwhals/stable/v2/__init__.py @@ -98,6 +98,8 @@ from narwhals.typing import ( IntoDType, IntoExpr, + IntoFrame, + IntoFrameT, IntoSchema, NonNestedLiteral, PythonLiteral, @@ -111,8 +113,10 @@ P = ParamSpec("P") R = TypeVar("R") +FrameT = TypeVar("FrameT", bound="IntoFrame") -class DataFrame(NwDataFrame[IntoDataFrameT]): + +class DataFrame(NwDataFrame[FrameT]): _version = Version.V2 @inherit_doc(NwDataFrame) @@ -240,7 +244,7 @@ def is_unique(self) -> Series[Any]: return _stableify(super().is_unique()) -class LazyFrame(NwLazyFrame[IntoLazyFrameT]): +class LazyFrame(NwLazyFrame[FrameT]): _version = Version.V2 @inherit_doc(NwLazyFrame) @@ -397,6 +401,10 @@ def from_native( @overload def from_native(native_object: LazyFrameT, **kwds: Unpack[AllowLazy]) -> LazyFrameT: ... @overload +def from_native( + native_object: LazyFrameT | DataFrameT, **kwds: Unpack[AllowLazy] +) -> LazyFrameT | DataFrameT: ... +@overload def from_native( native_object: IntoDataFrameT, **kwds: Unpack[ExcludeSeries] ) -> DataFrame[IntoDataFrameT]: ... @@ -413,10 +421,18 @@ def from_native( native_object: IntoLazyFrameT, **kwds: Unpack[AllowLazy] ) -> LazyFrame[IntoLazyFrameT]: ... @overload -def from_native( +def from_native( # type: ignore[overload-overlap] native_object: IntoDataFrameT | IntoSeriesT, **kwds: Unpack[AllowSeries] ) -> DataFrame[IntoDataFrameT] | Series[IntoSeriesT]: ... @overload +def from_native( + native_object: IntoDataFrameT | IntoLazyFrameT, **kwds: Unpack[AllowLazy] +) -> DataFrame[IntoDataFrameT] | LazyFrame[IntoLazyFrameT]: ... +@overload +def from_native( # type: ignore[overload-overlap] + native_object: IntoFrameT, **kwds: Unpack[AllowLazy] +) -> DataFrame[IntoFrameT] | LazyFrame[IntoFrameT]: ... +@overload def from_native( native_object: IntoDataFrameT | IntoLazyFrameT | IntoSeriesT, **kwds: Unpack[AllowAny] ) -> DataFrame[IntoDataFrameT] | LazyFrame[IntoLazyFrameT] | Series[IntoSeriesT]: ... @@ -433,7 +449,8 @@ def from_native( allow_series: bool | None, ) -> Any: ... def from_native( - native_object: IntoLazyFrameT + native_object: IntoFrameT + | IntoLazyFrameT | IntoDataFrameT | IntoSeriesT | IntoFrame @@ -444,7 +461,14 @@ def from_native( eager_only: bool = False, series_only: bool = False, allow_series: bool | None = None, -) -> LazyFrame[IntoLazyFrameT] | DataFrame[IntoDataFrameT] | Series[IntoSeriesT] | T: +) -> ( + LazyFrame[IntoLazyFrameT] + | DataFrame[IntoDataFrameT] + | LazyFrame[IntoFrameT] + | DataFrame[IntoFrameT] + | Series[IntoSeriesT] + | T +): """Convert `native_object` to Narwhals Dataframe, Lazyframe, or Series. Arguments: diff --git a/narwhals/translate.py b/narwhals/translate.py index 9452111bdd..f9ef32461b 100644 --- a/narwhals/translate.py +++ b/narwhals/translate.py @@ -55,6 +55,7 @@ Frame, IntoDataFrameT, IntoFrame, + IntoFrameT, IntoLazyFrameT, IntoSeries, IntoSeriesT, @@ -70,12 +71,12 @@ @overload def to_native( - narwhals_object: DataFrame[IntoDataFrameT], *, pass_through: Literal[False] = ... -) -> IntoDataFrameT: ... + narwhals_object: DataFrame[IntoFrameT], *, pass_through: Literal[False] = ... +) -> IntoFrameT: ... @overload def to_native( - narwhals_object: LazyFrame[IntoLazyFrameT], *, pass_through: Literal[False] = ... -) -> IntoLazyFrameT: ... + narwhals_object: LazyFrame[IntoFrameT], *, pass_through: Literal[False] = ... +) -> IntoFrameT: ... @overload def to_native( narwhals_object: Series[IntoSeriesT], *, pass_through: Literal[False] = ... @@ -85,12 +86,10 @@ def to_native(narwhals_object: Any, *, pass_through: bool) -> Any: ... def to_native( - narwhals_object: DataFrame[IntoDataFrameT] - | LazyFrame[IntoLazyFrameT] - | Series[IntoSeriesT], + narwhals_object: DataFrame[IntoFrameT] | LazyFrame[IntoFrameT] | Series[IntoSeriesT], *, pass_through: bool = False, -) -> IntoDataFrameT | IntoLazyFrameT | IntoSeriesT | Any: +) -> IntoFrameT | IntoSeriesT | Any: """Convert Narwhals object to native one. Arguments: @@ -129,6 +128,10 @@ def from_native( @overload def from_native(native_object: LazyFrameT, **kwds: Unpack[AllowLazy]) -> LazyFrameT: ... @overload +def from_native( + native_object: LazyFrameT | DataFrameT, **kwds: Unpack[AllowLazy] +) -> LazyFrameT | DataFrameT: ... +@overload def from_native( native_object: IntoDataFrameT, **kwds: Unpack[ExcludeSeries] ) -> DataFrame[IntoDataFrameT]: ... @@ -145,10 +148,18 @@ def from_native( native_object: IntoLazyFrameT, **kwds: Unpack[AllowLazy] ) -> LazyFrame[IntoLazyFrameT]: ... @overload -def from_native( +def from_native( # type: ignore[overload-overlap] native_object: IntoDataFrameT | IntoSeriesT, **kwds: Unpack[AllowSeries] ) -> DataFrame[IntoDataFrameT] | Series[IntoSeriesT]: ... @overload +def from_native( + native_object: IntoDataFrameT | IntoLazyFrameT, **kwds: Unpack[AllowLazy] +) -> DataFrame[IntoDataFrameT] | LazyFrame[IntoLazyFrameT]: ... +@overload +def from_native( # type: ignore[overload-overlap] + native_object: IntoFrameT, **kwds: Unpack[AllowLazy] +) -> DataFrame[IntoFrameT] | LazyFrame[IntoFrameT]: ... +@overload def from_native( native_object: IntoDataFrameT | IntoLazyFrameT | IntoSeriesT, **kwds: Unpack[AllowAny] ) -> DataFrame[IntoDataFrameT] | LazyFrame[IntoLazyFrameT] | Series[IntoSeriesT]: ... @@ -164,8 +175,9 @@ def from_native( series_only: bool, allow_series: bool | None, ) -> Any: ... -def from_native( - native_object: IntoLazyFrameT +def from_native( # type: ignore[misc] + native_object: IntoFrameT + | IntoLazyFrameT | IntoDataFrameT | IntoSeriesT | IntoFrame @@ -176,7 +188,14 @@ def from_native( eager_only: bool = False, series_only: bool = False, allow_series: bool | None = None, -) -> LazyFrame[IntoLazyFrameT] | DataFrame[IntoDataFrameT] | Series[IntoSeriesT] | T: +) -> ( + LazyFrame[IntoLazyFrameT] + | DataFrame[IntoDataFrameT] + | LazyFrame[IntoFrameT] + | DataFrame[IntoFrameT] + | Series[IntoSeriesT] + | T +): """Convert `native_object` to Narwhals Dataframe, Lazyframe, or Series. Arguments: diff --git a/tests/translate/from_native_test.py b/tests/translate/from_native_test.py index 8d076699c0..deb36ccd0c 100644 --- a/tests/translate/from_native_test.py +++ b/tests/translate/from_native_test.py @@ -255,7 +255,7 @@ def test_init_already_narwhals_stable_to_unstable() -> None: s = native["a"] stable_s = nw_v1.from_native(s, allow_series=True) # type: ignore[var-annotated] - unstablified_s = nw.from_native(stable_s, allow_series=True) # type: ignore[var-annotated] + unstablified_s = nw.from_native(stable_s, allow_series=True) assert isinstance(unstablified_s, nw.Series) diff --git a/tests/translate/into_frame_t_test.py b/tests/translate/into_frame_t_test.py new file mode 100644 index 0000000000..98876bd841 --- /dev/null +++ b/tests/translate/into_frame_t_test.py @@ -0,0 +1,75 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +import pytest + +import narwhals as nw + +if TYPE_CHECKING: + from typing_extensions import assert_type + + from narwhals.typing import IntoFrameT + + +def test_readme_example() -> None: + # check that readme example (as of March 2026) passes + def _agnostic_function( # pragma: no cover + df_native: IntoFrameT, date_column: str, price_column: str + ) -> IntoFrameT: + df = nw.from_native(df_native) + if TYPE_CHECKING: + assert_type(df, nw.DataFrame[IntoFrameT] | nw.LazyFrame[IntoFrameT]) + res = ( + df.group_by(nw.col(date_column).dt.truncate("1mo")) + .agg(nw.col(price_column).mean()) + .sort(date_column) + ) + if TYPE_CHECKING: + assert_type(res, nw.DataFrame[IntoFrameT] | nw.LazyFrame[IntoFrameT]) + native = res.to_native() + if TYPE_CHECKING: + assert_type(native, IntoFrameT) + return res.to_native() + + +def test_from_eager_or_lazy_polars() -> None: + pytest.importorskip("polars") + import polars as pl + + lf = pl.LazyFrame() + df = pl.DataFrame() + either = lf if df.height else df + + r_lf = nw.from_native(lf) + r_df = nw.from_native(df) + r_either = nw.from_native(either) + + r2_lf = nw.from_native(r_lf) + r2_df = nw.from_native(r_df) + r2_either = nw.from_native(r_either) + + if TYPE_CHECKING: + assert_type(r_lf, nw.LazyFrame[pl.LazyFrame]) + assert_type(r_df, nw.DataFrame[pl.DataFrame]) + assert_type(r_either, nw.DataFrame[pl.DataFrame] | nw.LazyFrame[pl.LazyFrame]) + assert_type(r2_lf, nw.LazyFrame[pl.LazyFrame]) + assert_type(r2_df, nw.DataFrame[pl.DataFrame]) + assert_type(r2_either, nw.DataFrame[pl.DataFrame] | nw.LazyFrame[pl.LazyFrame]) + + +def test_into_frame_t_incompatible_apis() -> None: + def _agnostic_function( # pragma: no cover + df_native: IntoFrameT, + ) -> IntoFrameT: + nw.from_native(df_native).sink_parquet("...") # type: ignore[union-attr] + _ = ( + nw.from_native(df_native) + .with_row_index() # type: ignore[call-arg] + .to_native() + ) + return ( + nw.from_native(df_native) + .unique(maintain_order=True) # type: ignore[call-arg] + .to_native() + ) diff --git a/tests/v2_test.py b/tests/v2_test.py index edb14f8d90..7fd8c009ee 100644 --- a/tests/v2_test.py +++ b/tests/v2_test.py @@ -30,7 +30,7 @@ from typing_extensions import assert_type from narwhals._typing import EagerAllowed - from narwhals.stable.v2.typing import IntoDataFrameT + from narwhals.stable.v2.typing import IntoDataFrameT, IntoFrameT from narwhals.typing import IntoDType, _1DArray, _2DArray @@ -546,3 +546,46 @@ def test_first_last() -> None: result = df.select(b=nw_v2.col("a").first(), c=nw_v2.col("a").last()) expected = {"b": [0], "c": [-1]} assert_equal_data(result, expected) + + +def test_readme_example() -> None: + # check that readme example (as of March 2026) passes + def _agnostic_function( # pragma: no cover + df_native: IntoFrameT, date_column: str, price_column: str + ) -> IntoFrameT: + return ( + nw_v2.from_native(df_native) + .group_by(nw_v2.col(date_column).dt.truncate("1mo")) + .agg(nw_v2.col(price_column).mean()) + .sort(date_column) + .to_native() + ) + + +def test_from_eager_or_lazy_polars() -> None: + pytest.importorskip("polars") + import polars as pl + + lf = pl.LazyFrame() + df = pl.DataFrame() + either = lf if df.height else df + + r_lf = nw_v2.from_native(lf) + r_df = nw_v2.from_native(df) + r_either = nw_v2.from_native(either) + + r2_lf = nw_v2.from_native(r_lf) + r2_df = nw_v2.from_native(r_df) + r2_either = nw_v2.from_native(r_either) + + if TYPE_CHECKING: + assert_type(r_lf, nw_v2.LazyFrame[pl.LazyFrame]) + assert_type(r_df, nw_v2.DataFrame[pl.DataFrame]) + assert_type( + r_either, nw_v2.DataFrame[pl.DataFrame] | nw_v2.LazyFrame[pl.LazyFrame] + ) + assert_type(r2_lf, nw_v2.LazyFrame[pl.LazyFrame]) + assert_type(r2_df, nw_v2.DataFrame[pl.DataFrame]) + assert_type( + r2_either, nw_v2.DataFrame[pl.DataFrame] | nw_v2.LazyFrame[pl.LazyFrame] + )