fix(typing): Align lambda_expr to cpp implementation

Author: dangotbanned

narwhals/_duckdb/typing.py

@@ -1,9 +1,10 @@
 from __future__ import annotations
 
 from collections.abc import Sequence
-from typing import TYPE_CHECKING, Any, Literal, Protocol, TypedDict, overload
+from typing import TYPE_CHECKING, Any, Literal, Protocol, TypedDict, Union, overload
 
 import duckdb
+from duckdb import Expression
 
 from narwhals._typing_compat import TypeVar
 
@@ -12,13 +13,19 @@
 
     import numpy as np
     import pandas as pd
-    from duckdb import DuckDBPyConnection, Expression
+    from duckdb import DuckDBPyConnection
     from typing_extensions import TypeAlias, TypeIs
 
     from narwhals.typing import Into1DArray, PythonLiteral
 
 
-__all__ = ["BaseType", "WindowExpressionKwargs", "has_children", "is_dtype"]
+__all__ = [
+    "BaseType",
+    "IntoColumnExpr",
+    "WindowExpressionKwargs",
+    "has_children",
+    "is_dtype",
+]
 
 IntoDuckDBLiteral: TypeAlias = """
     PythonLiteral
@@ -60,6 +67,8 @@ class WindowExpressionKwargs(TypedDict, total=False):
 _Enum: TypeAlias = Literal["enum"]
 _Decimal: TypeAlias = Literal["decimal"]
 _TimestampTZ: TypeAlias = Literal["timestamp with time zone"]
+IntoColumnExpr: TypeAlias = Union[str, Expression]
+"""A column name, or the result of calling `duckdb.ColumnExpression`."""
 
 
 class BaseType(Protocol[_ID_co]):

narwhals/_duckdb/utils.py

@@ -6,7 +6,13 @@
 import duckdb
 from duckdb import Expression
 
-from narwhals._duckdb.typing import BaseType, IntoDuckDBLiteral, has_children, is_dtype
+from narwhals._duckdb.typing import (
+    BaseType,
+    IntoColumnExpr,
+    IntoDuckDBLiteral,
+    has_children,
+    is_dtype,
+)
 from narwhals._utils import Implementation, Version, isinstance_or_issubclass, zip_strict
 from narwhals.exceptions import ColumnNotFoundError, UnsupportedDTypeError
 
@@ -54,6 +60,7 @@
 
 
 # TODO @dangotbanned: Raise an issue upstream on `Expression | str` too narrow
+# NOTE: https://github.com/duckdb/duckdb-python/blob/df7789cbd31b2d2b8d03d012f14331bc3297fb2d/src/duckdb_py/native/python_conversion.cpp#L916-L1069
 def lit(value: IntoDuckDBLiteral | Expression) -> Expression:
     """Alias for `duckdb.ConstantExpression`."""
     lit_: Incomplete = duckdb.ConstantExpression
@@ -67,10 +74,10 @@ def lit(value: IntoDuckDBLiteral | Expression) -> Expression:
 """Alias for `duckdb.FunctionExpression`."""
 
 
-# TODO @dangotbanned: Investigate `lhs: Expression | str | tuple[str]`
-# Seems incorrect
+# TODO @dangotbanned: Raise an issue upstream on `Expression | str | tuple[str` too narrow
+# NOTE: https://github.com/duckdb/duckdb-python/blob/df7789cbd31b2d2b8d03d012f14331bc3297fb2d/src/duckdb_py/pyexpression.cpp#L361-L413
 def lambda_expr(
-    params: str | Expression | tuple[Expression, ...], expr: Expression, /
+    params: IntoColumnExpr | tuple[IntoColumnExpr, ...], expr: Expression, /
 ) -> Expression:
     """Wraps [`duckdb.LambdaExpression`].