Skip to content

Commit 189dff4

Browse files
authored
refactor(typing): Simplify concat signature (#2339)
* refactor(typing): Simplify `concat` signature - Using a constrained `TypeVar`, like in `polars` - Discovered in #2337 * test(typing): Correctly ignore on invalid test Indicates the new signature is working * chore(typing): Just cast for now
1 parent d2a9f42 commit 189dff4

File tree

3 files changed

+9
-57
lines changed

3 files changed

+9
-57
lines changed

narwhals/functions.py

Lines changed: 4 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -9,8 +9,8 @@
99
from typing import Literal
1010
from typing import Mapping
1111
from typing import Sequence
12+
from typing import TypeVar
1213
from typing import cast
13-
from typing import overload
1414

1515
from narwhals._expression_parsing import ExpansionKind
1616
from narwhals._expression_parsing import ExprKind
@@ -57,46 +57,21 @@
5757
from narwhals.dtypes import DType
5858
from narwhals.schema import Schema
5959
from narwhals.series import Series
60-
from narwhals.typing import IntoDataFrameT
6160
from narwhals.typing import IntoExpr
62-
from narwhals.typing import IntoFrameT
6361
from narwhals.typing import IntoSeriesT
6462
from narwhals.typing import NativeFrame
6563
from narwhals.typing import NativeLazyFrame
6664
from narwhals.typing import _2DArray
6765

6866
_IntoSchema: TypeAlias = "Mapping[str, DType] | Schema | Sequence[str] | None"
67+
FrameT = TypeVar("FrameT", "DataFrame[Any]", "LazyFrame[Any]")
6968

7069

71-
@overload
7270
def concat(
73-
items: Iterable[DataFrame[IntoDataFrameT]],
71+
items: Iterable[FrameT],
7472
*,
7573
how: Literal["horizontal", "vertical", "diagonal"] = "vertical",
76-
) -> DataFrame[IntoDataFrameT]: ...
77-
78-
79-
@overload
80-
def concat(
81-
items: Iterable[LazyFrame[IntoFrameT]],
82-
*,
83-
how: Literal["horizontal", "vertical", "diagonal"] = "vertical",
84-
) -> LazyFrame[IntoFrameT]: ...
85-
86-
87-
@overload
88-
def concat(
89-
items: Iterable[DataFrame[IntoDataFrameT] | LazyFrame[IntoFrameT]],
90-
*,
91-
how: Literal["horizontal", "vertical", "diagonal"] = "vertical",
92-
) -> DataFrame[IntoDataFrameT] | LazyFrame[IntoFrameT]: ...
93-
94-
95-
def concat(
96-
items: Iterable[DataFrame[IntoDataFrameT] | LazyFrame[IntoFrameT]],
97-
*,
98-
how: Literal["horizontal", "vertical", "diagonal"] = "vertical",
99-
) -> DataFrame[IntoDataFrameT] | LazyFrame[IntoFrameT]:
74+
) -> FrameT:
10075
"""Concatenate multiple DataFrames, LazyFrames into a single entity.
10176
10277
Arguments:

narwhals/stable/v1/__init__.py

Lines changed: 4 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -99,6 +99,7 @@
9999
from narwhals.typing import _1DArray
100100
from narwhals.typing import _2DArray
101101

102+
FrameT = TypeVar("FrameT", "DataFrame[Any]", "LazyFrame[Any]")
102103
IntoSeriesT = TypeVar("IntoSeriesT", bound="IntoSeries", default=Any)
103104
T = TypeVar("T", default=Any)
104105
else:
@@ -2049,35 +2050,11 @@ def max_horizontal(*exprs: IntoExpr | Iterable[IntoExpr]) -> Expr:
20492050
return _stableify(nw.max_horizontal(*exprs))
20502051

20512052

2052-
@overload
2053-
def concat(
2054-
items: Iterable[DataFrame[IntoDataFrameT]],
2055-
*,
2056-
how: Literal["horizontal", "vertical", "diagonal"] = "vertical",
2057-
) -> DataFrame[IntoDataFrameT]: ...
2058-
2059-
2060-
@overload
2061-
def concat(
2062-
items: Iterable[LazyFrame[IntoFrameT]],
2063-
*,
2064-
how: Literal["horizontal", "vertical", "diagonal"] = "vertical",
2065-
) -> LazyFrame[IntoFrameT]: ...
2066-
2067-
2068-
@overload
2069-
def concat(
2070-
items: Iterable[DataFrame[IntoDataFrameT] | LazyFrame[IntoFrameT]],
2071-
*,
2072-
how: Literal["horizontal", "vertical", "diagonal"] = "vertical",
2073-
) -> DataFrame[IntoDataFrameT] | LazyFrame[IntoFrameT]: ...
2074-
2075-
20762053
def concat(
2077-
items: Iterable[DataFrame[IntoDataFrameT] | LazyFrame[IntoFrameT]],
2054+
items: Iterable[FrameT],
20782055
*,
20792056
how: Literal["horizontal", "vertical", "diagonal"] = "vertical",
2080-
) -> DataFrame[IntoDataFrameT] | LazyFrame[IntoFrameT]:
2057+
) -> FrameT:
20812058
"""Concatenate multiple DataFrames, LazyFrames into a single entity.
20822059
20832060
Arguments:
@@ -2096,7 +2073,7 @@ def concat(
20962073
Raises:
20972074
TypeError: The items to concatenate should either all be eager, or all lazy
20982075
"""
2099-
return _stableify(nw.concat(items, how=how))
2076+
return cast("FrameT", _stableify(nw.concat(items, how=how)))
21002077

21012078

21022079
def concat_str(

tests/frame/invalid_test.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -83,7 +83,7 @@ def test_validate_laziness() -> None:
8383
TypeError,
8484
match=("The items to concatenate should either all be eager, or all lazy"),
8585
):
86-
nw.concat([nw.from_native(df, eager_only=True), nw.from_native(df).lazy()])
86+
nw.concat([nw.from_native(df, eager_only=True), nw.from_native(df).lazy()]) # type: ignore[type-var]
8787

8888

8989
@pytest.mark.slow

0 commit comments

Comments
 (0)