use overload

narwhals-dev · Aug 25, 2024 · f60793e · f60793e
1 parent 84ba6c2
commit f60793e
Showing 1 changed file with 18 additions and 3 deletions.
diff --git a/narwhals/stable/v1.py b/narwhals/stable/v1.py
@@ -60,7 +60,6 @@
     from narwhals.typing import IntoExpr
 
 T = TypeVar("T")
-FrameT = TypeVar("FrameT", "DataFrame[Any]", "LazyFrame[Any]")
 
 
 class DataFrame(NwDataFrame[IntoDataFrameT]):
@@ -1338,11 +1337,27 @@ def mean_horizontal(*exprs: IntoExpr | Iterable[IntoExpr]) -> Expr:
     return _stableify(nw.mean_horizontal(*exprs))
 
 
+@overload
+def concat(
+    items: Iterable[NwDataFrame[Any]],
+    *,
+    how: Literal["horizontal", "vertical"] = "vertical",
+) -> DataFrame[Any]: ...
+
+
+@overload
+def concat(
+    items: Iterable[NwLazyFrame[Any]],
+    *,
+    how: Literal["horizontal", "vertical"] = "vertical",
+) -> LazyFrame[Any]: ...
+
+
 def concat(
-    items: Iterable[FrameT],
+    items: Iterable[NwDataFrame[Any] | NwLazyFrame[Any]],
     *,
     how: Literal["horizontal", "vertical"] = "vertical",
-) -> FrameT:
+) -> DataFrame[Any] | LazyFrame[Any]:
     return _stableify(nw.concat(items, how=how))  # type: ignore[no-any-return]