feat: Improve Literal type unwrapping

- 4th iteration of trying to get this working
- A nice property of this is the unimplemented `RangeLiteral` case is (statically) unreachable in `impl_arrow`
  - Without the need to name *any* of the `LiteralValue` classes
- Also the traversal is now hidden behind `unwrap`, which preserves the type

Author: dangotbanned

narwhals/_plan/expr.py

@@ -13,13 +13,14 @@
 import typing as t
 
 from narwhals._plan.aggregation import Agg, OrderableAgg
-from narwhals._plan.common import ExprIR, SelectorIR, _field_str
+from narwhals._plan.common import ExprIR, SelectorIR, _field_str, is_non_nested_literal
 from narwhals._plan.name import KeepName, RenameAlias
 from narwhals._plan.typing import (
     ExprT,
     FunctionT,
     LeftSelectorT,
     LeftT,
+    LiteralT,
     Ns,
     OperatorT,
     RightSelectorT,
@@ -115,12 +116,12 @@ def to_compliant(self, plx: Ns[ExprT], /) -> ExprT:
         return plx.col(*self.names)
 
 
-class Literal(ExprIR):
+class Literal(ExprIR, t.Generic[LiteralT]):
     """https://github.com/pola-rs/polars/blob/dafd0a2d0e32b52bcfa4273bffdd6071a0d5977a/crates/polars-plan/src/dsl/expr.rs#L81."""
 
     __slots__ = ("value",)
 
-    value: LiteralValue
+    value: LiteralValue[LiteralT]
 
     @property
     def is_scalar(self) -> bool:
@@ -138,7 +139,13 @@ def __repr__(self) -> str:
         return f"lit({self.value!r})"
 
     def to_compliant(self, plx: Ns[ExprT], /) -> ExprT:
-        return plx.lit(self.value.unwrap(), self.dtype)
+        value = self.unwrap()
+        if is_non_nested_literal(value):
+            return plx.lit(value, self.dtype)
+        raise NotImplementedError(type(self.value))
+
+    def unwrap(self) -> LiteralT:
+        return self.value.unwrap()
 
 
 class _BinaryOp(ExprIR, t.Generic[LeftT, OperatorT, RightT]):

narwhals/_plan/impl_arrow.py

@@ -9,13 +9,15 @@
 from functools import singledispatch
 
 from narwhals._plan import expr
-from narwhals._plan.literal import is_scalar_literal, is_series_literal
+from narwhals._plan.literal import is_literal_scalar, is_literal_series
 
 if t.TYPE_CHECKING:
     import pyarrow as pa
     from typing_extensions import TypeAlias
 
     from narwhals._plan.common import ExprIR
+    from narwhals._plan.dummy import DummySeries
+    from narwhals.typing import NonNestedLiteral
 
     NativeFrame: TypeAlias = pa.Table
     NativeSeries: TypeAlias = pa.ChunkedArray[t.Any]
@@ -38,15 +40,17 @@ def cols(node: expr.Columns, frame: NativeFrame) -> Evaluated:
 
 
 @evaluate.register(expr.Literal)
-def lit(node: expr.Literal, frame: NativeFrame) -> Evaluated:
+def lit(
+    node: expr.Literal[NonNestedLiteral] | expr.Literal[DummySeries], frame: NativeFrame
+) -> Evaluated:
     import pyarrow as pa
 
-    if is_scalar_literal(node.value):
+    if is_literal_scalar(node):
         lit: t.Any = pa.scalar
-        array = pa.repeat(lit(node.value.unwrap()), len(frame))
+        array = pa.repeat(lit(node.unwrap()), len(frame))
         return [pa.chunked_array([array])]
-    elif is_series_literal(node.value):
-        ca = node.value.unwrap().to_native()
+    elif is_literal_series(node):
+        ca = node.unwrap().to_native()
         return [t.cast("NativeSeries", ca)]
     else:
         raise NotImplementedError(type(node.value))

narwhals/_plan/literal.py

@@ -3,24 +3,17 @@
 from typing import TYPE_CHECKING, Any, Generic
 
 from narwhals._plan.common import Immutable
+from narwhals._plan.typing import LiteralT, NonNestedLiteralT
 
 if TYPE_CHECKING:
     from typing_extensions import TypeIs
 
     from narwhals._plan.dummy import DummySeries
     from narwhals._plan.expr import Literal
     from narwhals.dtypes import DType
-    from narwhals.typing import NonNestedLiteral
 
-from narwhals._typing_compat import TypeVar
 
-T = TypeVar("T", default=Any)
-NonNestedLiteralT = TypeVar(
-    "NonNestedLiteralT", bound="NonNestedLiteral", default="NonNestedLiteral"
-)
-
-
-class LiteralValue(Immutable, Generic[T]):
+class LiteralValue(Immutable, Generic[LiteralT]):
     """https://github.com/pola-rs/polars/blob/dafd0a2d0e32b52bcfa4273bffdd6071a0d5977a/crates/polars-plan/src/plans/lit.rs#L67-L73."""
 
     @property
@@ -40,7 +33,7 @@ def to_literal(self) -> Literal:
 
         return Literal(value=self)
 
-    def unwrap(self) -> T:
+    def unwrap(self) -> LiteralT:
         raise NotImplementedError
 
 
@@ -102,9 +95,27 @@ class RangeLiteral(LiteralValue):
     dtype: DType
 
 
-def is_scalar_literal(obj: Any) -> TypeIs[ScalarLiteral]:
+def _is_scalar(
+    obj: ScalarLiteral[NonNestedLiteralT] | Any,
+) -> TypeIs[ScalarLiteral[NonNestedLiteralT]]:
     return isinstance(obj, ScalarLiteral)
 
 
-def is_series_literal(obj: Any) -> TypeIs[SeriesLiteral]:
+def _is_series(obj: Any) -> TypeIs[SeriesLiteral]:
     return isinstance(obj, SeriesLiteral)
+
+
+def is_literal(obj: Literal[LiteralT] | Any) -> TypeIs[Literal[LiteralT]]:
+    from narwhals._plan.expr import Literal
+
+    return isinstance(obj, Literal)
+
+
+def is_literal_scalar(
+    obj: Literal[NonNestedLiteralT] | Any,
+) -> TypeIs[Literal[NonNestedLiteralT]]:
+    return is_literal(obj) and _is_scalar(obj.value)
+
+
+def is_literal_series(obj: Literal[DummySeries] | Any) -> TypeIs[Literal[DummySeries]]:
+    return is_literal(obj) and _is_series(obj.value)

narwhals/_plan/typing.py

@@ -11,7 +11,9 @@
     from narwhals._compliant.typing import CompliantExprAny
     from narwhals._plan import operators as ops
     from narwhals._plan.common import ExprIR, Function, IRNamespace, SelectorIR
+    from narwhals._plan.dummy import DummySeries
     from narwhals._plan.functions import RollingWindow
+    from narwhals.typing import NonNestedLiteral
 
 __all__ = ["FunctionT", "LeftT", "OperatorT", "RightT", "RollingT", "SelectorOperatorT"]
 
@@ -29,6 +31,12 @@
     "SelectorOperatorT", bound="ops.SelectorOperator", default="ops.SelectorOperator"
 )
 IRNamespaceT = TypeVar("IRNamespaceT", bound="IRNamespace")
+
+NonNestedLiteralT = TypeVar(
+    "NonNestedLiteralT", bound="NonNestedLiteral", default="NonNestedLiteral"
+)
+LiteralT = TypeVar("LiteralT", bound="NonNestedLiteral | DummySeries", default=t.Any)
+
 # NOTE: Shorter aliases of `_compliant.typing`
 # - Aiming to try and preserve the types as much as possible
 # - Recursion between `Expr` and `Frame` is an issue