Repository: narwhals-dev/narwhals
Pull request #365 (2024-06-30 21:06:05)

feat: add IntoDataframe type hint
<!--
# Thanks for contributing a pull request! 
## Please make sure you see our contribution guidelines: https://github.com/narwhals-dev/narwhals/blob/main/CONTRIBUTING.md
-->

## What type of PR is this? (check all applicable)

- [ ] 💾 Refactor
- [ ] ✨ Feature
- [ ] 🐛 Bug Fix
- [ ] 🔧 Optimization
- [ ] 📝 Documentation
- [ ] ✅ Test
- [ ] 🐳 Other

## Related issues 

- Related issue # 
- Closes #

## Checklist

- [ ] Code follows style guide (ruff)
- [ ] Tests added 
- [ ] Documented the changes

## If you have comments or can explain your changes, please do so below.




Patch:
diff --git a/narwhals/translate.py b/narwhals/translate.py
index 37fd9f2b..729d4804 100644
--- a/narwhals/translate.py
+++ b/narwhals/translate.py
@@ -5,6 +5,7 @@ from typing import TYPE_CHECKING
 from typing import Any
 from typing import Callable
 from typing import Literal
+from typing import TypeVar
 from typing import overload
 
 from narwhals.dependencies import get_cudf
@@ -17,6 +18,9 @@ if TYPE_CHECKING:
     from narwhals.dataframe import DataFrame
     from narwhals.dataframe import LazyFrame
     from narwhals.series import Series
+    from narwhals.typing import IntoDataFrame
+
+T = TypeVar("T")
 
 
 def to_native(narwhals_object: Any, *, strict: bool = True) -> Any:
@@ -58,18 +62,73 @@ def to_native(narwhals_object: Any, *, strict: bool = True) -> Any:
 def from_native(
     native_dataframe: Any,
     *,
-    strict: bool = ...,
+    strict: Literal[False],
     eager_only: Literal[True],
     series_only: None = ...,
     allow_series: Literal[True],
-) -> DataFrame | Series: ...
+) -> Any: ...
+
+
+@overload
+def from_native(
+    native_dataframe: IntoDataFrame | T,
+    *,
+    strict: Literal[False],
+    eager_only: Literal[True],
+    series_only: None = ...,
+    allow_series: None = ...,
+) -> DataFrame | T: ...
+
+
+@overload
+def from_native(
+    native_dataframe: Any,
+    *,
+    strict: Literal[False],
+    eager_only: None = ...,
+    series_only: None = ...,
+    allow_series: Literal[True],
+) -> Any: ...
 
 
 @overload
 def from_native(
     native_dataframe: Any,
     *,
-    strict: bool = ...,
+    strict: Literal[False],
+    eager_only: None = ...,
+    series_only: Literal[True],
+    allow_series: None = ...,
+) -> Any: ...
+
+
+@overload
+def from_native(
+    native_dataframe: IntoDataFrame | T,
+    *,
+    strict: Literal[False],
+    eager_only: None = ...,
+    series_only: None = ...,
+    allow_series: None = ...,
+) -> DataFrame | LazyFrame | T: ...
+
+
+@overload
+def from_native(
+    native_dataframe: Any,
+    *,
+    strict: Literal[True] = ...,
+    eager_only: Literal[True],
+    series_only: None = ...,
+    allow_series: Literal[True],
+) -> DataFrame | Series: ...
+
+
+@overload
+def from_native(
+    native_dataframe: IntoDataFrame,
+    *,
+    strict: Literal[True] = ...,
     eager_only: Literal[True],
     series_only: None = ...,
     allow_series: None = ...,
@@ -80,7 +139,7 @@ def from_native(
 def from_native(
     native_dataframe: Any,
     *,
-    strict: bool = ...,
+    strict: Literal[True] = ...,
     eager_only: None = ...,
     series_only: None = ...,
     allow_series: Literal[True],
@@ -91,7 +150,7 @@ def from_native(
 def from_native(
     native_dataframe: Any,
     *,
-    strict: bool = ...,
+    strict: Literal[True] = ...,
     eager_only: None = ...,
     series_only: Literal[True],
     allow_series: None = ...,
@@ -100,15 +159,16 @@ def from_native(
 
 @overload
 def from_native(
-    native_dataframe: Any,
+    native_dataframe: IntoDataFrame,
     *,
-    strict: bool = ...,
+    strict: Literal[True] = ...,
     eager_only: None = ...,
     series_only: None = ...,
     allow_series: None = ...,
 ) -> DataFrame | LazyFrame: ...
 
 
+# Nothing was specified
 @overload
 def from_native(
     native_dataframe: Any,
@@ -117,7 +177,7 @@ def from_native(
     eager_only: bool | None,
     series_only: bool | None,
     allow_series: bool | None,
-) -> DataFrame | LazyFrame | Series: ...
+) -> Any: ...
 
 
 def from_native(
@@ -127,7 +187,7 @@ def from_native(
     eager_only: bool | None = None,
     series_only: bool | None = None,
     allow_series: bool | None = None,
-) -> DataFrame | LazyFrame | Series:
+) -> Any:
     """
     Convert dataframe to Narwhals DataFrame, LazyFrame, or Series.
 
@@ -220,7 +280,7 @@ def from_native(
     elif strict:  # pragma: no cover
         msg = f"Expected pandas-like dataframe, Polars dataframe, or Polars lazyframe, got: {type(native_dataframe)}"
         raise TypeError(msg)
-    return native_dataframe  # type: ignore[no-any-return]  # pragma: no cover (todo)
+    return native_dataframe  # pragma: no cover (todo)
 
 
 def get_native_namespace(obj: Any) -> Any:
diff --git a/narwhals/typing.py b/narwhals/typing.py
index d94a168f..d8a00c92 100644
--- a/narwhals/typing.py
+++ b/narwhals/typing.py
@@ -1,6 +1,9 @@
-from typing import TYPE_CHECKING  # pragma: no cover
-from typing import TypeVar  # pragma: no cover
-from typing import Union  # pragma: no cover
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+from typing import Any
+from typing import Protocol
+from typing import Union
 
 if TYPE_CHECKING:
     import sys
@@ -10,10 +13,23 @@ if TYPE_CHECKING:
     else:
         from typing_extensions import TypeAlias
 
+    from narwhals.dataframe import DataFrame
     from narwhals.expression import Expr
     from narwhals.series import Series
 
-    IntoExpr: TypeAlias = Union[Expr, str, int, float, Series]
+    # All dataframes supported by Narwhals have a
+    # `columns` property. Their similarities don't extend
+    # _that_ much further unfortunately...
+    class NativeDataFrame(Protocol):
+        @property
+        def columns(self) -> Any: ...
+
+        def join(self, *args: Any, **kwargs: Any) -> Any: ...
+
+
+# Anything which can be converted to an expression.
+IntoExpr: TypeAlias = Union["Expr", str, int, float, "Series"]
+# Anything which can be converted to a Narwhals DataFrame.
+IntoDataFrame: TypeAlias = Union["NativeDataFrame", "DataFrame"]
 
-    NativeDataFrame = TypeVar("NativeDataFrame")
-    NativeSeries = TypeVar("NativeSeries")
+__all__ = ["IntoExpr", "IntoDataFrame"]
diff --git a/pyproject.toml b/pyproject.toml
index 3dfd82c0..80039b83 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -85,6 +85,7 @@ markers = ["slow: marks tests as slow (deselect with '-m \"not slow\"')"]
 plugins = ["covdefaults"]
 
 [tool.coverage.report]
+omit = ['narwhals/typing.py']
 exclude_also = [
   "> POLARS_VERSION",
   "if sys.version_info() <",


Test patch:
diff --git a/tests/expr/str/starts_with_ends_with_test.py b/tests/expr/str/starts_with_ends_with_test.py
index 6ce69e7e..13c7ddf6 100644
--- a/tests/expr/str/starts_with_ends_with_test.py
+++ b/tests/expr/str/starts_with_ends_with_test.py
@@ -7,6 +7,10 @@ import polars as pl
 import pytest
 
 import narwhals as nw
+
+# Don't move this into typechecking block, for coverage
+# purposes
+from narwhals.typing import IntoDataFrame  # noqa: TCH001
 from tests.utils import compare_dicts
 from tests.utils import maybe_get_modin_df
 
@@ -19,7 +23,7 @@ df_mpd = maybe_get_modin_df(df_pandas)
     "df_raw",
     [df_pandas, df_lazy, df_mpd],
 )
-def test_ends_with(df_raw: Any) -> None:
+def test_ends_with(df_raw: IntoDataFrame) -> None:
     df = nw.from_native(df_raw).lazy()
     result = df.select(nw.col("a").str.ends_with("das"))
     result_native = nw.to_native(result)