better typing

Author: MarcoGorelli

narwhals/_expression_parsing.py

@@ -18,8 +18,6 @@
 from narwhals.utils import Implementation
 
 if TYPE_CHECKING:
-    from typing_extensions import TypeAlias
-
     from narwhals._arrow.expr import ArrowExpr
     from narwhals._pandas_like.expr import PandasLikeExpr
     from narwhals.typing import CompliantDataFrame
@@ -28,13 +26,9 @@
     from narwhals.typing import CompliantNamespace
     from narwhals.typing import CompliantSeries
     from narwhals.typing import CompliantSeriesT_co
+    from narwhals.typing import IntoCompliantExpr
     from narwhals.typing import IntoExpr
 
-    IntoCompliantExpr: TypeAlias = (
-        CompliantExpr[CompliantSeriesT_co] | str | CompliantSeriesT_co
-    )
-    CompliantExprT = TypeVar("CompliantExprT", bound=CompliantExpr[Any])
-
     ArrowOrPandasLikeExpr = TypeVar(
         "ArrowOrPandasLikeExpr", bound=Union[ArrowExpr, PandasLikeExpr]
     )

narwhals/_polars/dataframe.py

@@ -214,11 +214,7 @@ def __getitem__(self: Self, item: Any) -> Any:
             return self._from_native_object(result)
 
     def simple_select(self, *column_names: str) -> Self:
-        try:
-            return self._from_native_frame(self._native_frame.select(*column_names))
-        except Exception as e:
-            msg = f"{e!s}\n\nHint: Did you mean one of these columns: {self.columns}?"
-            raise ColumnNotFoundError(msg) from e
+        return self._from_native_frame(self._native_frame.select(*column_names))
 
     def get_column(self: Self, name: str) -> PolarsSeries:
         from narwhals._polars.series import PolarsSeries
@@ -484,7 +480,4 @@ def unpivot(
         )
 
     def simple_select(self, *column_names: str) -> Self:
-        try:
-            return self._from_native_frame(self._native_frame.select(*column_names))
-        except Exception as e:
-            raise ColumnNotFoundError(str(e)) from e
+        return self._from_native_frame(self._native_frame.select(*column_names))

narwhals/dataframe.py

@@ -44,6 +44,7 @@
     from narwhals.group_by import GroupBy
     from narwhals.group_by import LazyGroupBy
     from narwhals.series import Series
+    from narwhals.typing import IntoCompliantExpr
     from narwhals.typing import IntoDataFrame
     from narwhals.typing import IntoExpr
     from narwhals.typing import IntoFrame
@@ -70,28 +71,29 @@ def __narwhals_namespace__(self: Self) -> Any:
     def _from_compliant_dataframe(self: Self, df: Any) -> Self:
         # construct, preserving properties
         return self.__class__(  # type: ignore[call-arg]
-            df,
-            level=self._level,
+            df, level=self._level
         )
 
     def _flatten_parse_col_names_into_expr_and_extract(
         self, *exprs: IntoExpr | Any, **named_exprs: IntoExpr | Any
-    ) -> Any:
+    ) -> tuple[tuple[IntoCompliantExpr[Any]], dict[str, IntoCompliantExpr[Any]]]:
         """Process `args` and `kwargs`, extracting underlying objects as we go, interpreting strings as column names."""
         plx = self.__narwhals_namespace__()
-        exprs = tuple(
+        compliant_exprs = tuple(
             plx.col(expr) if isinstance(expr, str) else self._extract_compliant(expr)
             for expr in flatten(exprs)
         )
-        named_exprs = {
+        compliant_named_exprs = {
             key: plx.col(value)
             if isinstance(value, str)
             else self._extract_compliant(value)
             for key, value in named_exprs.items()
         }
-        return exprs, named_exprs
+        return compliant_exprs, compliant_named_exprs
 
-    def _flatten_and_extract(self: Self, *args: Any, **kwargs: Any) -> Any:
+    def _flatten_and_extract(
+        self: Self, *args: Any, **kwargs: Any
+    ) -> tuple[tuple[IntoCompliantExpr[Any]], dict[str, IntoCompliantExpr[Any]]]:
         """Process `args` and `kwargs`, extracting underlying objects as we go."""
         args = [self._extract_compliant(v) for v in flatten(args)]  # type: ignore[assignment]
         kwargs = {k: self._extract_compliant(v) for k, v in kwargs.items()}
@@ -136,35 +138,35 @@ def columns(self: Self) -> list[str]:
     def with_columns(
         self: Self, *exprs: IntoExpr | Iterable[IntoExpr], **named_exprs: IntoExpr
     ) -> Self:
-        exprs, named_exprs = self._flatten_parse_col_names_into_expr_and_extract(
-            *exprs, **named_exprs
+        compliant_exprs, compliant_named_exprs = (
+            self._flatten_parse_col_names_into_expr_and_extract(*exprs, **named_exprs)
         )
         return self._from_compliant_dataframe(
-            self._compliant_frame.with_columns(*exprs, **named_exprs),
+            self._compliant_frame.with_columns(*compliant_exprs, **compliant_named_exprs),
         )
 
     def select(
         self: Self,
         *exprs: IntoExpr | Iterable[IntoExpr],
         **named_exprs: IntoExpr,
     ) -> Self:
-        flat_exprs = list(flatten(exprs))
-        if flat_exprs and all(isinstance(x, str) for x in flat_exprs) and not named_exprs:
+        exprs = tuple(flatten(exprs))
+        if exprs and all(isinstance(x, str) for x in exprs) and not named_exprs:
             # fast path!
             try:
                 return self._from_compliant_dataframe(
-                    self._compliant_frame.simple_select(*flat_exprs),
+                    self._compliant_frame.simple_select(*exprs),
                 )
             except Exception as e:
                 # Column not found is the only thing that can realistically be raised here.
                 msg = f"{e!s}\n\nHint: Did you mean one of these columns: {self.columns}?"
                 raise ColumnNotFoundError(msg) from e
 
-        flat_exprs, named_exprs = self._flatten_parse_col_names_into_expr_and_extract(
-            *flat_exprs, **named_exprs
+        compliant_exprs, compliant_named_exprs = (
+            self._flatten_parse_col_names_into_expr_and_extract(*exprs, **named_exprs)
         )
         return self._from_compliant_dataframe(
-            self._compliant_frame.select(*flat_exprs, **named_exprs),
+            self._compliant_frame.select(*compliant_exprs, **compliant_named_exprs),
         )
 
     def rename(self: Self, mapping: dict[str, str]) -> Self:

narwhals/group_by.py

@@ -122,9 +122,11 @@ def agg(
                 "but `df.group_by('a').agg(nw.col('b'))` is not."
             )
             raise InvalidOperationError(msg)
-        aggs, named_aggs = self._df._flatten_and_extract(*aggs, **named_aggs)
+        compliant_aggs, compliant_named_aggs = self._df._flatten_and_extract(
+            *aggs, **named_aggs
+        )
         return self._df._from_compliant_dataframe(  # type: ignore[return-value]
-            self._grouped.agg(*aggs, **named_aggs),
+            self._grouped.agg(*compliant_aggs, **compliant_named_aggs),
         )
 
     def __iter__(self: Self) -> Iterator[tuple[Any, DataFrameT]]:
@@ -218,7 +220,9 @@ def agg(
                 "but `df.group_by('a').agg(nw.col('b'))` is not."
             )
             raise InvalidOperationError(msg)
-        aggs, named_aggs = self._df._flatten_and_extract(*aggs, **named_aggs)
+        compliant_aggs, compliant_named_aggs = self._df._flatten_and_extract(
+            *aggs, **named_aggs
+        )
         return self._df._from_compliant_dataframe(  # type: ignore[return-value]
-            self._grouped.agg(*aggs, **named_aggs),
+            self._grouped.agg(*compliant_aggs, **compliant_named_aggs),
         )

narwhals/typing.py

@@ -95,6 +95,11 @@ def __mod__(self, other: Any) -> Self: ...
     def __pow__(self, other: Any) -> Self: ...
 
 
+IntoCompliantExpr: TypeAlias = (
+    CompliantExpr[CompliantSeriesT_co] | str | CompliantSeriesT_co
+)
+
+
 class CompliantNamespace(Protocol, Generic[CompliantSeriesT_co]):
     def col(self, *column_names: str) -> CompliantExpr[CompliantSeriesT_co]: ...
     def lit(

tests/expr_and_series/concat_str_test.py

@@ -38,7 +38,7 @@ def test_concat_str(
             nw.concat_str(
                 [
                     nw.col("a") * 2,
-                    nw.col("b"),
+                    "b",
                     nw.col("c"),
                 ],
                 separator=" ",