feat: Infer DType in lit

Author: dangotbanned

narwhals/_plan/common.py

@@ -22,6 +22,7 @@
     from narwhals._plan.lists import IRListNamespace
     from narwhals._plan.meta import IRMetaNamespace
     from narwhals._plan.options import FunctionOptions
+    from narwhals.dtypes import DType
     from narwhals.typing import NonNestedLiteral
 
 else:
@@ -315,3 +316,21 @@ def is_iterable_reject(obj: Any) -> TypeIs[str | bytes | DummySeries]:
     from narwhals._plan.dummy import DummySeries
 
     return isinstance(obj, (str, bytes, DummySeries))
+
+
+def py_to_narwhals_dtype(obj: NonNestedLiteral, version: Version = Version.MAIN) -> DType:
+    dtypes = version.dtypes
+    mapping: dict[type[NonNestedLiteral], type[DType]] = {
+        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)()

narwhals/_plan/demo.py

@@ -9,7 +9,13 @@
     expr_parsing as parse,
     functions as F,  # noqa: N812
 )
-from narwhals._plan.common import ExprIR, IntoExpr, is_expr, is_non_nested_literal
+from narwhals._plan.common import (
+    ExprIR,
+    IntoExpr,
+    is_expr,
+    is_non_nested_literal,
+    py_to_narwhals_dtype,
+)
 from narwhals._plan.dummy import DummySeries
 from narwhals._plan.expr import All, Column, Columns, IndexColumns, Len, Nth
 from narwhals._plan.literal import ScalarLiteral, SeriesLiteral
@@ -52,11 +58,11 @@ def lit(
 ) -> DummyExpr:
     if isinstance(value, DummySeries):
         return SeriesLiteral(value=value).to_literal().to_narwhals()
-    if dtype is None or not isinstance(dtype, DType):
-        dtype = Version.MAIN.dtypes.Unknown()
     if not is_non_nested_literal(value):
         msg = f"{type(value).__name__!r} is not supported in `nw.lit`, got: {value!r}."
         raise TypeError(msg)
+    if dtype is None or not isinstance(dtype, DType):
+        dtype = py_to_narwhals_dtype(value, Version.MAIN)
     return ScalarLiteral(value=value, dtype=dtype).to_literal().to_narwhals()