fixup flattening

Author: MarcoGorelli

narwhals/dataframe.py

@@ -72,7 +72,7 @@ def _from_dataframe(self, df: Any) -> Self:
     def _flatten_and_extract(self, *args: Any, **kwargs: Any) -> Any:
         from narwhals.utils import flatten
 
-        args = [self._extract_native(v) for v in flatten(*args)]  # type: ignore[assignment]
+        args = [self._extract_native(v) for v in flatten(args)]  # type: ignore[assignment]
         kwargs = {k: self._extract_native(v) for k, v in kwargs.items()}
         return args, kwargs
 

narwhals/expression.py

@@ -70,6 +70,16 @@ def __rsub__(self, other: Any) -> Expr:
             lambda plx: self._call(plx).__rsub__(extract_native(plx, other))
         )
 
+    def __truediv__(self, other: Any) -> Expr:
+        return self.__class__(
+            lambda plx: self._call(plx).__truediv__(extract_native(plx, other))
+        )
+
+    def __rtruediv__(self, other: Any) -> Expr:
+        return self.__class__(
+            lambda plx: self._call(plx).__rtruediv__(extract_native(plx, other))
+        )
+
     def __mul__(self, other: Any) -> Expr:
         return self.__class__(
             lambda plx: self._call(plx).__mul__(extract_native(plx, other))

narwhals/pandas_like/utils.py

@@ -89,7 +89,7 @@ def parse_into_exprs(
     *exprs: IntoPandasExpr | Iterable[IntoPandasExpr],
     **named_exprs: IntoPandasExpr,
 ) -> list[PandasExpr]:
-    out = [parse_into_expr(implementation, into_expr) for into_expr in flatten(*exprs)]  # type: ignore[arg-type]
+    out = [parse_into_expr(implementation, into_expr) for into_expr in flatten(exprs)]
     for name, expr in named_exprs.items():
         out.append(parse_into_expr(implementation, expr).alias(name))
     return out
@@ -128,8 +128,7 @@ def evaluate_into_exprs(
     series: list[PandasSeries] = [
         item
         for sublist in [
-            evaluate_into_expr(df, into_expr)  # type: ignore[arg-type]
-            for into_expr in flatten(*exprs)
+            evaluate_into_expr(df, into_expr) for into_expr in flatten(*exprs)
         ]
         for item in sublist
     ]

narwhals/series.py

@@ -130,3 +130,9 @@ def to_numpy(self) -> Any:
 
     def to_pandas(self) -> Any:
         return self._series.to_pandas()
+
+    def mean(self) -> Any:
+        return self._series.mean()
+
+    def std(self) -> Any:
+        return self._series.std()

narwhals/utils.py

@@ -1,26 +1,27 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+from typing import Any
 from typing import Iterable
 
-if TYPE_CHECKING:
-    from narwhals.typing import T
-
 
 def remove_prefix(text: str, prefix: str) -> str:
     if text.startswith(prefix):
         return text[len(prefix) :]
     return text
 
 
-def flatten(*args: T | Iterable[T]) -> list[T]:
-    out: list[T] = []
-    for arg in args:
-        if isinstance(arg, (list, tuple)):
-            out.extend(arg)
-        else:
-            out.append(arg)  # type: ignore[arg-type]
-    return out
+def flatten(args: Any) -> list[Any]:
+    if not args:
+        return []
+    if len(args) == 1 and _is_iterable(args[0]):
+        return args[0]  # type: ignore[no-any-return]
+    return args  # type: ignore[no-any-return]
+
+
+def _is_iterable(arg: Any | Iterable[Any]) -> bool:
+    from narwhals.series import Series
+
+    return isinstance(arg, Iterable) and not isinstance(arg, (str, bytes, Series))
 
 
 def flatten_str(*args: str | Iterable[str]) -> list[str]:

t.py

@@ -2,11 +2,12 @@
 # type: ignore
 import polars
 import pandas as pd
+import polars as pl
 
 import narwhals as nw
 
 df_raw = pd.DataFrame({"a": [1, 3, 2], "b": [4, 4, 6], "z": [7.0, 8, 9]})
-df = nw.DataFrame(df_raw)
+df = nw.DataFrame(df_raw, features=["lazy"])
 df_raw_2 = pd.DataFrame({"a": [1, 3], "c": [7, 9]})
 df2 = nw.DataFrame(df_raw_2, features=["lazy"])
 
@@ -123,3 +124,8 @@
 
 result = df.select("a", "b").select(nw.all() + nw.col("a"))
 print(nw.to_native(result))
+
+df = nw.DataFrame(df_raw, features=["eager"])
+print(df["a"].mean())
+df = nw.DataFrame(pl.from_pandas(df_raw), features=["eager"])
+print(df["a"].mean())