feat: Support nested data in lit

Playing catch up with #3424
@FBruzzesi

The `list` stuff was the main wart left in `_parse.py`.
Was easier implement the `list` support in `lit`,
than try to make sense of whatever logic I landed on before 😅

Up next, is to fix the leaky `DataFrame.filter` list parsing!

https://github.com/narwhals-dev/narwhals/blob/6e43487472a13000d99361e7ccc3fcb8cc3e656f/narwhals/_plan/dataframe.py#L571-L579

Author: dangotbanned

narwhals/_plan/_guards.py

@@ -34,7 +34,7 @@
     bytes,
     Decimal,
 )
-_PYTHON_LITERAL_TPS = (*_NON_NESTED_LITERAL_TPS, list, tuple, type(None))
+_PYTHON_LITERAL_TPS = (*_NON_NESTED_LITERAL_TPS, list, tuple, dict, type(None))
 
 
 def _ir(*_: Any):  # type: ignore[no-untyped-def]  # noqa: ANN202
@@ -57,6 +57,10 @@ def is_python_literal(obj: Any) -> TypeIs[PythonLiteral]:
     return isinstance(obj, _PYTHON_LITERAL_TPS)
 
 
+def is_python_literal_type(tp: type[Any]) -> TypeIs[type[PythonLiteral]]:
+    return tp in _PYTHON_LITERAL_TPS
+
+
 def is_series(obj: Series[NativeSeriesT] | Any) -> TypeIs[Series[NativeSeriesT]]:
     return isinstance(obj, _series().Series)
 

narwhals/_plan/_parse.py

@@ -63,61 +63,9 @@
     "sort_by_into_seq_of_expr_ir",
 ]
 
-_RaisesInvalidIntoExprError: TypeAlias = "Any"
-"""
-Placeholder for multiple `Iterable[IntoExpr]`.
-
-We only support cases `a`, `b`, but the typing for most contexts is more permissive:
-
->>> import polars as pl
->>> df = pl.DataFrame({"one": ["A", "B", "A"], "two": [1, 2, 3], "three": [4, 5, 6]})
->>> a = ("one", "two")
->>> b = (["one", "two"],)
->>>
->>> c = ("one", ["two"])
->>> d = (["one"], "two")
->>> [df.select(*into) for into in (a, b, c, d)]
-[shape: (3, 2)
- ┌─────┬─────┐
- │ one ┆ two │
- │ --- ┆ --- │
- │ str ┆ i64 │
- ╞═════╪═════╡
- │ A   ┆ 1   │
- │ B   ┆ 2   │
- │ A   ┆ 3   │
- └─────┴─────┘,
- shape: (3, 2)
- ┌─────┬─────┐
- │ one ┆ two │
- │ --- ┆ --- │
- │ str ┆ i64 │
- ╞═════╪═════╡
- │ A   ┆ 1   │
- │ B   ┆ 2   │
- │ A   ┆ 3   │
- └─────┴─────┘,
- shape: (3, 2)
- ┌─────┬───────────┐
- │ one ┆ literal   │
- │ --- ┆ ---       │
- │ str ┆ list[str] │
- ╞═════╪═══════════╡
- │ A   ┆ ["two"]   │
- │ B   ┆ ["two"]   │
- │ A   ┆ ["two"]   │
- └─────┴───────────┘,
- shape: (3, 2)
- ┌───────────┬─────┐
- │ literal   ┆ two │
- │ ---       ┆ --- │
- │ list[str] ┆ i64 │
- ╞═══════════╪═════╡
- │ ["one"]   ┆ 1   │
- │ ["one"]   ┆ 2   │
- │ ["one"]   ┆ 3   │
- └───────────┴─────┘]
-"""
+# TODO @dangotbanned: Simplify the `list` special-casing, now that it is supported
+Incomplete: TypeAlias = "Any"
+"""Artifact from previous `lit(list)` rejection"""
 
 
 def into_expr_ir(
@@ -146,7 +94,7 @@ def into_expr_ir(
 
 def into_seq_of_expr_ir(
     first_input: OneOrIterable[IntoExpr] = (),
-    *more_inputs: IntoExpr | _RaisesInvalidIntoExprError,
+    *more_inputs: IntoExpr | Incomplete,
     **named_inputs: IntoExpr,
 ) -> Seq[ExprIR]:
     """Parse variadic inputs into a flat sequence of expressions."""
@@ -157,7 +105,7 @@ def into_seq_of_expr_ir(
 
 def predicates_constraints_into_expr_ir(
     first_predicate: OneOrIterable[IntoExprColumn] | list[bool] = (),
-    *more_predicates: IntoExprColumn | list[bool] | _RaisesInvalidIntoExprError,
+    *more_predicates: IntoExprColumn | list[bool] | Incomplete,
     _list_as_series: PartialSeries | None = None,
     **constraints: IntoExpr,
 ) -> ExprIR:

narwhals/_plan/arrow/namespace.py

@@ -117,7 +117,7 @@ def col(self, node: ir.Column, frame: Frame, name: str) -> Expr:
             frame.native.column(node.name), name, version=frame.version
         )
 
-    def lit(self, node: ir.Lit[NonNestedLiteral], frame: Frame, name: str) -> Scalar:
+    def lit(self, node: ir.Lit[PythonLiteral], frame: Frame, name: str) -> Scalar:
         return self._scalar.from_python(
             node.value, name, dtype=node.dtype, version=frame.version
         )

narwhals/_plan/common.py

@@ -1,10 +1,8 @@
 from __future__ import annotations
 
-import datetime as dt
 import sys
 from collections.abc import Iterable
 from copy import deepcopy
-from decimal import Decimal
 from io import BytesIO
 from secrets import token_hex
 from types import MappingProxyType
@@ -14,7 +12,6 @@
 from narwhals._utils import _hasattr_static, qualified_type_name
 from narwhals.dtypes import DType
 from narwhals.exceptions import NarwhalsError
-from narwhals.utils import Version
 
 if TYPE_CHECKING:
     import reprlib
@@ -34,7 +31,7 @@
         Seq,
     )
     from narwhals._utils import _StoresColumns
-    from narwhals.typing import FileSource, NonNestedDType, NonNestedLiteral
+    from narwhals.typing import FileSource
 
     T = TypeVar("T")
 
@@ -52,24 +49,6 @@ def replace(obj: T, /, **changes: Any) -> T:
         return func(obj, **changes)  # type: ignore[no-any-return]
 
 
-def py_to_narwhals_dtype(obj: NonNestedLiteral, version: Version = Version.MAIN) -> DType:
-    dtypes = version.dtypes
-    mapping: dict[type[NonNestedLiteral], type[NonNestedDType]] = {
-        int: dtypes.Int64,
-        float: dtypes.Float64,
-        str: dtypes.String,
-        bool: dtypes.Boolean,
-        dt.datetime: dtypes.Datetime,
-        dt.date: dtypes.Date,
-        dt.time: dtypes.Time,
-        dt.timedelta: dtypes.Duration,
-        bytes: dtypes.Binary,
-        Decimal: dtypes.Decimal,
-        type(None): dtypes.Unknown,
-    }
-    return mapping.get(type(obj), dtypes.Unknown)()
-
-
 @overload
 def into_dtype(dtype: type[NonNestedDTypeT], /) -> NonNestedDTypeT: ...
 @overload

narwhals/_plan/compliant/namespace.py

@@ -23,7 +23,7 @@
     from narwhals._plan.expressions import FunctionExpr, boolean, functions as F
     from narwhals._plan.expressions.strings import ConcatStr
     from narwhals._utils import Implementation
-    from narwhals.typing import NonNestedLiteral
+    from narwhals.typing import PythonLiteral
 
 Incomplete: TypeAlias = Any
 
@@ -61,7 +61,7 @@ def coalesce(
     ) -> ExprT_co | ScalarT_co: ...
     def len(self, node: ir.Len, frame: FrameT, name: str) -> ScalarT_co: ...
     def lit(
-        self, node: ir.Lit[NonNestedLiteral], frame: FrameT, name: str
+        self, node: ir.Lit[PythonLiteral], frame: FrameT, name: str
     ) -> ScalarT_co: ...
     def max_horizontal(
         self, node: FunctionExpr[F.MaxHorizontal], frame: FrameT, name: str

narwhals/_plan/exceptions.py

@@ -69,8 +69,10 @@ def function_arg_non_scalar_error(
     return ShapeError(msg)
 
 
-def list_literal_error(value: Any) -> TypeError:
-    msg = f"{type(value).__name__!r} is not supported in `nw.lit`, got: {value!r}."
+def literal_type_error(value: Any) -> TypeError:
+    msg = f"{qualified_type_name(value)!r} is not supported in `nw.lit`"
+    if not isinstance(value, type):
+        msg = f"{msg}, got: {value!r}."
     return TypeError(msg)
 
 

narwhals/_plan/expressions/literal.py

@@ -2,27 +2,33 @@
 # see `LiteralExpr.value`
 from __future__ import annotations
 
+import datetime as dt
+from decimal import Decimal
+from functools import cache
 from typing import TYPE_CHECKING, Any, Generic, final
 
 from narwhals._plan import common
 from narwhals._plan._dispatch import DispatcherOptions
 from narwhals._plan._dtype import ResolveDType
 from narwhals._plan._expr_ir import ExprIR
+from narwhals._plan._guards import is_python_literal_type
+from narwhals._plan.exceptions import literal_type_error
 from narwhals._plan.typing import (
     LiteralT_co,
     NativeSeriesT,
     NativeSeriesT_co,
-    NonNestedLiteralT,
-    NonNestedLiteralT_co,
+    PythonLiteralT,
+    PythonLiteralT_co,
 )
+from narwhals._utils import Version
+from narwhals.dtypes import Field, List, Struct, Unknown
 
 if TYPE_CHECKING:
     from collections.abc import Iterator
 
     from narwhals._plan.series import Series
-    from narwhals._utils import Version
     from narwhals.dtypes import DType
-    from narwhals.typing import IntoDType
+    from narwhals.typing import IntoDType, NonNestedDType, PythonLiteral
 
 __all__ = ["Lit", "LitSeries", "lit", "lit_series"]
 
@@ -58,7 +64,7 @@ def iter_output_name(self) -> Iterator[ExprIR]:
 
 
 @final
-class Lit(LiteralExpr[NonNestedLiteralT_co], dispatch=namespaced()):
+class Lit(LiteralExpr[PythonLiteralT_co], dispatch=namespaced()):
     """An expression representing a scalar literal value.
 
     >>> import narwhals._plan as nw
@@ -78,6 +84,12 @@ def __repr__(self) -> str:
         v = self.value
         return f"lit({'null' if v is None else f'{type(v).__name__}: {v!s}'})"
 
+    @property
+    def __immutable_values__(self) -> Iterator[Any]:
+        dtype = self.dtype
+        value: Any = self.value
+        yield from (id(value) if dtype.is_nested() else value, dtype)
+
 
 @final
 class LitSeries(LiteralExpr["Series[NativeSeriesT_co]"], dispatch=namespaced()):
@@ -128,14 +140,73 @@ def __immutable_values__(self) -> Iterator[Any]:
         yield from (self.name, self.dtype, id(self.value))
 
 
-def lit(
-    value: NonNestedLiteralT, dtype: IntoDType | None = None
-) -> Lit[NonNestedLiteralT]:
-    if dtype is None:
-        dtype = common.py_to_narwhals_dtype(value)
-    else:
-        dtype = common.into_dtype(dtype)
+lit_series = LitSeries.from_series
+
+
+def lit(value: PythonLiteralT, dtype: IntoDType | None = None) -> Lit[PythonLiteralT]:
+    dtype = (
+        _py_value_to_dtype(value, Version.MAIN, allow_null=True)
+        if dtype is None
+        else common.into_dtype(dtype)
+    )
     return Lit(value=value, dtype=dtype)
 
 
-lit_series = LitSeries.from_series
+def _py_value_to_dtype(
+    obj: PythonLiteral, version: Version = Version.MAIN, *, allow_null: bool
+) -> DType:
+    # NOTE: Surely mypy must have fixed `_lru_cache_wrapper` hashable in a new version?
+    if dtype := _py_type_to_dtype(type(obj), version):  # type: ignore[arg-type]
+        if allow_null or not isinstance(dtype, Unknown):
+            return dtype
+        msg = "Nested dtypes containing nulls are not yet supported"
+        raise TypeError(msg)
+    if not isinstance(obj, (list, dict, tuple)):
+        # Just a type narrowing issue
+        msg = f"Expected unreachable, got {obj!r}"
+        raise NotImplementedError(msg)
+
+    if not obj:
+        msg = "Cannot infer dtype for empty nested structure. Please provide an explicit dtype parameter."
+        raise TypeError(msg)
+    if not isinstance(obj, dict):
+        first_value = next((el for el in obj if el is not None), None)
+        return List(_py_value_to_dtype(first_value, version, allow_null=False))
+    return Struct(
+        [
+            Field(k, _py_value_to_dtype(v, version, allow_null=False))
+            for k, v in obj.items()
+        ]
+    )
+
+
+@cache
+def _py_type_to_dtype(
+    py_type: type[PythonLiteral], version: Version = Version.MAIN, /
+) -> NonNestedDType | None:
+    """SAFETY.
+
+    Cache size is bound by these dimensions:
+
+        n_valid_py_types = len(non_nested) + len([list, dict, tuple])
+        maxsize = n_valid_py_types * len(Version)
+    """
+    dtypes = version.dtypes
+    non_nested: dict[type[PythonLiteral], type[NonNestedDType]] = {
+        int: dtypes.Int64,
+        float: dtypes.Float64,
+        str: dtypes.String,
+        bool: dtypes.Boolean,
+        dt.datetime: dtypes.Datetime,
+        dt.date: dtypes.Date,
+        dt.time: dtypes.Time,
+        dt.timedelta: dtypes.Duration,
+        bytes: dtypes.Binary,
+        Decimal: dtypes.Decimal,
+        type(None): dtypes.Unknown,
+    }
+    if dtype := non_nested.get(py_type):
+        return dtype()
+    if not is_python_literal_type(py_type):
+        raise literal_type_error(py_type)
+    return None

narwhals/_plan/functions/literal.py

@@ -3,20 +3,20 @@
 from typing import TYPE_CHECKING
 
 from narwhals._plan import _guards, expressions as ir
-from narwhals._plan.exceptions import list_literal_error
+from narwhals._plan.exceptions import literal_type_error
 
 if TYPE_CHECKING:
     from narwhals._plan.expr import Expr
     from narwhals._plan.series import Series
     from narwhals._plan.typing import NativeSeriesT
-    from narwhals.typing import IntoDType, NonNestedLiteral
+    from narwhals.typing import IntoDType, PythonLiteral
 
 
 def lit(
-    value: NonNestedLiteral | Series[NativeSeriesT], dtype: IntoDType | None = None
+    value: PythonLiteral | Series[NativeSeriesT], dtype: IntoDType | None = None
 ) -> Expr:
     if _guards.is_series(value):
         return ir.lit_series(value).to_narwhals()
-    if not _guards.is_non_nested_literal(value):
-        raise list_literal_error(value)
+    if not _guards.is_python_literal(value):
+        raise literal_type_error(value)
     return ir.lit(value, dtype).to_narwhals()

narwhals/_plan/polars/namespace.py

@@ -35,7 +35,7 @@
         ConcatMethod,
         IntoDType,
         IntoSchema,
-        NonNestedLiteral,
+        PythonLiteral,
     )
 
 Incomplete: TypeAlias = Any
@@ -169,7 +169,7 @@ def col(self, node: ir.Column, frame: Incomplete, name: str) -> Expr:
     def len(self, node: ir.Len, frame: Incomplete, name: str) -> Expr:
         return self._expr.from_native(pl.len(), name, self.version)
 
-    def lit(self, node: ir.Lit[NonNestedLiteral], frame: Incomplete, name: str) -> Expr:
+    def lit(self, node: ir.Lit[PythonLiteral], frame: Incomplete, name: str) -> Expr:
         return self._expr.from_python(
             node.value, name, dtype=node.dtype, version=self.version
         )