Fix Expression typing in relational API stubs

Author: evertlammerts

_duckdb-stubs/__init__.pyi

@@ -1,6 +1,9 @@
+import datetime
+import decimal
 import os
 import pathlib
 import typing
+import uuid
 from typing_extensions import Self
 
 if typing.TYPE_CHECKING:
@@ -15,6 +18,22 @@ if typing.TYPE_CHECKING:
     # the field_ids argument to to_parquet and write_parquet has a recursive structure
     ParquetFieldIdsType = Mapping[str, typing.Union[int, "ParquetFieldIdsType"]]
 
+_ExpressionLike: typing.TypeAlias = typing.Union[
+    "Expression",
+    str,
+    int,
+    float,
+    bool,
+    bytes,
+    None,
+    datetime.date,
+    datetime.datetime,
+    datetime.time,
+    datetime.timedelta,
+    decimal.Decimal,
+    uuid.UUID,
+]
+
 __all__: list[str] = [
     "BinderException",
     "CSVLineTerminator",
@@ -472,7 +491,7 @@ class DuckDBPyRelation:
     def __getitem__(self, name: str) -> DuckDBPyRelation: ...
     def __len__(self) -> int: ...
     def aggregate(
-        self, aggr_expr: str | Iterable[Expression | str], group_expr: Expression | str = ""
+        self, aggr_expr: str | Iterable[_ExpressionLike], group_expr: _ExpressionLike = ""
     ) -> DuckDBPyRelation: ...
     def any_value(
         self, expression: str, groups: str = "", window_spec: str = "", projected_columns: str = ""
@@ -642,7 +661,7 @@ class DuckDBPyRelation:
     def product(
         self, expression: str, groups: str = "", window_spec: str = "", projected_columns: str = ""
     ) -> DuckDBPyRelation: ...
-    def project(self, *args: str | Expression, groups: str = "") -> DuckDBPyRelation: ...
+    def project(self, *args: _ExpressionLike, groups: str = "") -> DuckDBPyRelation: ...
     def quantile(
         self,
         expression: str,
@@ -671,7 +690,7 @@ class DuckDBPyRelation:
     def rank(self, window_spec: str, projected_columns: str = "") -> DuckDBPyRelation: ...
     def rank_dense(self, window_spec: str, projected_columns: str = "") -> DuckDBPyRelation: ...
     def row_number(self, window_spec: str, projected_columns: str = "") -> DuckDBPyRelation: ...
-    def select(self, *args: str | Expression, groups: str = "") -> DuckDBPyRelation: ...
+    def select(self, *args: _ExpressionLike, groups: str = "") -> DuckDBPyRelation: ...
     def select_dtypes(self, types: typing.List[sqltypes.DuckDBPyType | str]) -> DuckDBPyRelation: ...
     def select_types(self, types: typing.List[sqltypes.DuckDBPyType | str]) -> DuckDBPyRelation: ...
     def set_alias(self, alias: str) -> DuckDBPyRelation: ...
@@ -684,7 +703,7 @@ class DuckDBPyRelation:
         null_value: str | None = None,
         render_mode: RenderMode | None = None,
     ) -> None: ...
-    def sort(self, *args: Expression) -> DuckDBPyRelation: ...
+    def sort(self, *args: _ExpressionLike) -> DuckDBPyRelation: ...
     def sql_query(self) -> str: ...
     def std(
         self, expression: str, groups: str = "", window_spec: str = "", projected_columns: str = ""
@@ -748,7 +767,7 @@ class DuckDBPyRelation:
     def torch(self) -> dict[str, typing.Any]: ...
     def union(self, union_rel: DuckDBPyRelation) -> DuckDBPyRelation: ...
     def unique(self, unique_aggr: str) -> DuckDBPyRelation: ...
-    def update(self, set: Expression | str, *, condition: Expression | str | None = None) -> None: ...
+    def update(self, set: dict[str, _ExpressionLike], *, condition: _ExpressionLike | None = None) -> None: ...
     def value_counts(self, expression: str, groups: str = "") -> DuckDBPyRelation: ...
     def var(
         self, expression: str, groups: str = "", window_spec: str = "", projected_columns: str = ""
@@ -856,54 +875,54 @@ class ExplainType:
     def value(self) -> int: ...
 
 class Expression:
-    def __add__(self, other: Expression) -> Expression: ...
-    def __and__(self, other: Expression) -> Expression: ...
-    def __div__(self, other: Expression) -> Expression: ...
-    def __eq__(self, other: Expression) -> Expression: ...  # type: ignore[override]
-    def __floordiv__(self, other: Expression) -> Expression: ...
-    def __ge__(self, other: Expression) -> Expression: ...
-    def __gt__(self, other: Expression) -> Expression: ...
+    def __add__(self, other: _ExpressionLike) -> Expression: ...
+    def __and__(self, other: _ExpressionLike) -> Expression: ...
+    def __div__(self, other: _ExpressionLike) -> Expression: ...
+    def __eq__(self, other: _ExpressionLike) -> Expression: ...  # type: ignore[override]
+    def __floordiv__(self, other: _ExpressionLike) -> Expression: ...
+    def __ge__(self, other: _ExpressionLike) -> Expression: ...
+    def __gt__(self, other: _ExpressionLike) -> Expression: ...
     @typing.overload
     def __init__(self, arg0: str) -> None: ...
     @typing.overload
     def __init__(self, arg0: typing.Any) -> None: ...
     def __invert__(self) -> Expression: ...
-    def __le__(self, other: Expression) -> Expression: ...
-    def __lt__(self, other: Expression) -> Expression: ...
-    def __mod__(self, other: Expression) -> Expression: ...
-    def __mul__(self, other: Expression) -> Expression: ...
-    def __ne__(self, other: Expression) -> Expression: ...  # type: ignore[override]
+    def __le__(self, other: _ExpressionLike) -> Expression: ...
+    def __lt__(self, other: _ExpressionLike) -> Expression: ...
+    def __mod__(self, other: _ExpressionLike) -> Expression: ...
+    def __mul__(self, other: _ExpressionLike) -> Expression: ...
+    def __ne__(self, other: _ExpressionLike) -> Expression: ...  # type: ignore[override]
     def __neg__(self) -> Expression: ...
-    def __or__(self, other: Expression) -> Expression: ...
-    def __pow__(self, other: Expression) -> Expression: ...
-    def __radd__(self, other: Expression) -> Expression: ...
-    def __rand__(self, other: Expression) -> Expression: ...
-    def __rdiv__(self, other: Expression) -> Expression: ...
-    def __rfloordiv__(self, other: Expression) -> Expression: ...
-    def __rmod__(self, other: Expression) -> Expression: ...
-    def __rmul__(self, other: Expression) -> Expression: ...
-    def __ror__(self, other: Expression) -> Expression: ...
-    def __rpow__(self, other: Expression) -> Expression: ...
-    def __rsub__(self, other: Expression) -> Expression: ...
-    def __rtruediv__(self, other: Expression) -> Expression: ...
-    def __sub__(self, other: Expression) -> Expression: ...
-    def __truediv__(self, other: Expression) -> Expression: ...
+    def __or__(self, other: _ExpressionLike) -> Expression: ...
+    def __pow__(self, other: _ExpressionLike) -> Expression: ...
+    def __radd__(self, other: _ExpressionLike) -> Expression: ...
+    def __rand__(self, other: _ExpressionLike) -> Expression: ...
+    def __rdiv__(self, other: _ExpressionLike) -> Expression: ...
+    def __rfloordiv__(self, other: _ExpressionLike) -> Expression: ...
+    def __rmod__(self, other: _ExpressionLike) -> Expression: ...
+    def __rmul__(self, other: _ExpressionLike) -> Expression: ...
+    def __ror__(self, other: _ExpressionLike) -> Expression: ...
+    def __rpow__(self, other: _ExpressionLike) -> Expression: ...
+    def __rsub__(self, other: _ExpressionLike) -> Expression: ...
+    def __rtruediv__(self, other: _ExpressionLike) -> Expression: ...
+    def __sub__(self, other: _ExpressionLike) -> Expression: ...
+    def __truediv__(self, other: _ExpressionLike) -> Expression: ...
     def alias(self, name: str) -> Expression: ...
     def asc(self) -> Expression: ...
-    def between(self, lower: Expression, upper: Expression) -> Expression: ...
+    def between(self, lower: _ExpressionLike, upper: _ExpressionLike) -> Expression: ...
     def cast(self, type: sqltypes.DuckDBPyType) -> Expression: ...
     def collate(self, collation: str) -> Expression: ...
     def desc(self) -> Expression: ...
     def get_name(self) -> str: ...
-    def isin(self, *args: Expression) -> Expression: ...
-    def isnotin(self, *args: Expression) -> Expression: ...
+    def isin(self, *args: _ExpressionLike) -> Expression: ...
+    def isnotin(self, *args: _ExpressionLike) -> Expression: ...
     def isnotnull(self) -> Expression: ...
     def isnull(self) -> Expression: ...
     def nulls_first(self) -> Expression: ...
     def nulls_last(self) -> Expression: ...
-    def otherwise(self, value: Expression) -> Expression: ...
+    def otherwise(self, value: _ExpressionLike) -> Expression: ...
     def show(self) -> None: ...
-    def when(self, condition: Expression, value: Expression) -> Expression: ...
+    def when(self, condition: _ExpressionLike, value: _ExpressionLike) -> Expression: ...
 
 class FatalException(DatabaseError): ...
 
@@ -1055,18 +1074,18 @@ class token_type:
     @property
     def value(self) -> int: ...
 
-def CaseExpression(condition: Expression, value: Expression) -> Expression: ...
-def CoalesceOperator(*args: Expression) -> Expression: ...
+def CaseExpression(condition: _ExpressionLike, value: _ExpressionLike) -> Expression: ...
+def CoalesceOperator(*args: _ExpressionLike) -> Expression: ...
 def ColumnExpression(*args: str) -> Expression: ...
 def ConstantExpression(value: typing.Any) -> Expression: ...
 def DefaultExpression() -> Expression: ...
-def FunctionExpression(function_name: str, *args: Expression) -> Expression: ...
-def LambdaExpression(lhs: typing.Any, rhs: Expression) -> Expression: ...
+def FunctionExpression(function_name: str, *args: _ExpressionLike) -> Expression: ...
+def LambdaExpression(lhs: typing.Any, rhs: _ExpressionLike) -> Expression: ...
 def SQLExpression(expression: str) -> Expression: ...
 def StarExpression(*, exclude: Iterable[str | Expression] | None = None) -> Expression: ...
 def aggregate(
     df: pandas.DataFrame,
-    aggr_expr: str | Iterable[Expression | str],
+    aggr_expr: str | Iterable[_ExpressionLike],
     group_expr: str = "",
     *,
     connection: DuckDBPyConnection | None = None,
@@ -1324,7 +1343,7 @@ def pl(
     connection: DuckDBPyConnection | None = None,
 ) -> typing.Union[polars.DataFrame, polars.LazyFrame]: ...
 def project(
-    df: pandas.DataFrame, *args: str | Expression, groups: str = "", connection: DuckDBPyConnection | None = None
+    df: pandas.DataFrame, *args: _ExpressionLike, groups: str = "", connection: DuckDBPyConnection | None = None
 ) -> DuckDBPyRelation: ...
 def query(
     query: Statement | str,