refactor of datatypes typing:

- fixed all places where it was too narrow. Most of the time str are accepted for sqltypes. odd exception seems to be the map method on Relation
- using Self for annotations on arguments when pertinent

Author: OutSquareCapital

_duckdb-stubs/__init__.pyi

@@ -12,7 +12,15 @@ if typing.TYPE_CHECKING:
     import pyarrow.lib
     from builtins import list as lst
     from collections.abc import Callable, Iterable, Sequence, Mapping
-    from ._typing import ParquetFieldIdsType, IntoExpr, IntoExprColumn, PythonLiteral, IntoValues
+    from ._typing import (
+        ParquetFieldIdsType,
+        IntoExpr,
+        IntoExprColumn,
+        PythonLiteral,
+        IntoValues,
+        IntoDType,
+        IntoNestedDType,
+    )
     from duckdb import sqltypes, func
 
 __all__: lst[str] = [
@@ -193,7 +201,7 @@ class DuckDBPyConnection:
     def __enter__(self) -> Self: ...
     def __exit__(self, exc_type: object, exc: object, traceback: object) -> None: ...
     def append(self, table_name: str, df: pandas.DataFrame, *, by_name: bool = False) -> DuckDBPyConnection: ...
-    def array_type(self, type: sqltypes.DuckDBPyType, size: typing.SupportsInt) -> sqltypes.DuckDBPyType: ...
+    def array_type(self, type: IntoDType, size: typing.SupportsInt) -> sqltypes.DuckDBPyType: ...
     def arrow(self, rows_per_batch: typing.SupportsInt = 1000000) -> pyarrow.lib.RecordBatchReader:
         """Alias of to_arrow_reader(). We recommend using to_arrow_reader() instead."""
         ...
@@ -207,8 +215,8 @@ class DuckDBPyConnection:
         self,
         name: str,
         function: Callable[..., typing.Any],
-        parameters: lst[sqltypes.DuckDBPyType] | None = None,
-        return_type: sqltypes.DuckDBPyType | None = None,
+        parameters: lst[IntoDType] | None = None,
+        return_type: IntoDType | None = None,
         *,
         type: func.PythonUDFType = ...,
         null_handling: func.FunctionNullHandling = ...,
@@ -327,9 +335,9 @@ class DuckDBPyConnection:
     def disable_profiling(self) -> None: ...
     def interrupt(self) -> None: ...
     def list_filesystems(self) -> lst[str]: ...
-    def list_type(self, type: sqltypes.DuckDBPyType) -> sqltypes.DuckDBPyType: ...
+    def list_type(self, type: IntoDType) -> sqltypes.DuckDBPyType: ...
     def load_extension(self, extension: str) -> None: ...
-    def map_type(self, key: sqltypes.DuckDBPyType, value: sqltypes.DuckDBPyType) -> sqltypes.DuckDBPyType: ...
+    def map_type(self, key: IntoDType, value: IntoDType) -> sqltypes.DuckDBPyType: ...
     @typing.overload
     def pl(
         self, rows_per_batch: typing.SupportsInt = 1000000, *, lazy: typing.Literal[False] = ...
@@ -439,23 +447,17 @@ class DuckDBPyConnection:
     def register_filesystem(self, filesystem: fsspec.AbstractFileSystem) -> None: ...
     def remove_function(self, name: str) -> DuckDBPyConnection: ...
     def rollback(self) -> DuckDBPyConnection: ...
-    def row_type(
-        self, fields: dict[str, sqltypes.DuckDBPyType] | lst[sqltypes.DuckDBPyType]
-    ) -> sqltypes.DuckDBPyType: ...
+    def row_type(self, fields: IntoNestedDType) -> sqltypes.DuckDBPyType: ...
     def sql(self, query: Statement | str, *, alias: str = "", params: object = None) -> DuckDBPyRelation: ...
     def sqltype(self, type_str: str) -> sqltypes.DuckDBPyType: ...
     def string_type(self, collation: str = "") -> sqltypes.DuckDBPyType: ...
-    def struct_type(
-        self, fields: dict[str, sqltypes.DuckDBPyType] | lst[sqltypes.DuckDBPyType]
-    ) -> sqltypes.DuckDBPyType: ...
+    def struct_type(self, fields: IntoNestedDType) -> sqltypes.DuckDBPyType: ...
     def table(self, table_name: str) -> DuckDBPyRelation: ...
     def table_function(self, name: str, parameters: object = None) -> DuckDBPyRelation: ...
     def tf(self) -> dict[str, typing.Any]: ...
     def torch(self) -> dict[str, typing.Any]: ...
     def type(self, type_str: str) -> sqltypes.DuckDBPyType: ...
-    def union_type(
-        self, members: lst[sqltypes.DuckDBPyType] | dict[str, sqltypes.DuckDBPyType]
-    ) -> sqltypes.DuckDBPyType: ...
+    def union_type(self, members: IntoNestedDType) -> sqltypes.DuckDBPyType: ...
     def unregister(self, view_name: str) -> DuckDBPyConnection: ...
     def unregister_filesystem(self, name: str) -> None: ...
     def values(self, *args: IntoValues) -> DuckDBPyRelation: ...
@@ -527,13 +529,13 @@ class DuckDBPyRelation:
     ) -> DuckDBPyRelation: ...
     def create(self, table_name: str) -> None: ...
     def create_view(self, view_name: str, replace: bool = True) -> DuckDBPyRelation: ...
-    def cross(self, other_rel: DuckDBPyRelation) -> DuckDBPyRelation: ...
+    def cross(self, other_rel: Self) -> DuckDBPyRelation: ...
     def cume_dist(self, window_spec: str, projected_columns: str = "") -> DuckDBPyRelation: ...
     def dense_rank(self, window_spec: str, projected_columns: str = "") -> DuckDBPyRelation: ...
     def describe(self) -> DuckDBPyRelation: ...
     def df(self, *, date_as_object: bool = False) -> pandas.DataFrame: ...
     def distinct(self) -> DuckDBPyRelation: ...
-    def except_(self, other_rel: DuckDBPyRelation) -> DuckDBPyRelation: ...
+    def except_(self, other_rel: Self) -> DuckDBPyRelation: ...
     def execute(self) -> DuckDBPyRelation: ...
     def explain(self, type: ExplainType = ExplainType.STANDARD) -> str: ...
     def favg(
@@ -568,10 +570,10 @@ class DuckDBPyRelation:
     ) -> DuckDBPyRelation: ...
     def insert(self, values: lst[object]) -> None: ...
     def insert_into(self, table_name: str) -> None: ...
-    def intersect(self, other_rel: DuckDBPyRelation) -> DuckDBPyRelation: ...
+    def intersect(self, other_rel: Self) -> DuckDBPyRelation: ...
     def join(
         self,
-        other_rel: DuckDBPyRelation,
+        other_rel: Self,
         condition: IntoExprColumn,
         how: typing.Literal["inner", "left", "right", "outer", "semi", "anti"] = "inner",
     ) -> DuckDBPyRelation: ...
@@ -747,7 +749,7 @@ class DuckDBPyRelation:
     def to_table(self, table_name: str) -> None: ...
     def to_view(self, view_name: str, replace: bool = True) -> DuckDBPyRelation: ...
     def torch(self) -> dict[str, typing.Any]: ...
-    def union(self, union_rel: DuckDBPyRelation) -> DuckDBPyRelation: ...
+    def union(self, union_rel: Self) -> DuckDBPyRelation: ...
     def unique(self, unique_aggr: str) -> DuckDBPyRelation: ...
     def update(self, set: Mapping[str, IntoExpr], *, condition: IntoExpr = None) -> None: ...
     def value_counts(self, expression: str, groups: str = "") -> DuckDBPyRelation: ...
@@ -1027,7 +1029,7 @@ def append(
     table_name: str, df: pandas.DataFrame, *, by_name: bool = False, connection: DuckDBPyConnection | None = None
 ) -> DuckDBPyConnection: ...
 def array_type(
-    type: sqltypes.DuckDBPyType, size: typing.SupportsInt, *, connection: DuckDBPyConnection | None = None
+    type: IntoDType, size: typing.SupportsInt, *, connection: DuckDBPyConnection | None = None
 ) -> sqltypes.DuckDBPyType: ...
 @typing.overload
 def arrow(
@@ -1056,8 +1058,8 @@ def connect(
 def create_function(
     name: str,
     function: Callable[..., typing.Any],
-    parameters: lst[sqltypes.DuckDBPyType] | None = None,
-    return_type: sqltypes.DuckDBPyType | None = None,
+    parameters: lst[IntoDType] | None = None,
+    return_type: IntoDType | None = None,
     *,
     type: func.PythonUDFType = ...,
     null_handling: func.FunctionNullHandling = ...,
@@ -1240,15 +1242,10 @@ def get_profiling_information(*, connection: DuckDBPyConnection | None = None, f
 def enable_profiling(*, connection: DuckDBPyConnection | None = None) -> None: ...
 def disable_profiling(*, connection: DuckDBPyConnection | None = None) -> None: ...
 def list_filesystems(*, connection: DuckDBPyConnection | None = None) -> lst[str]: ...
-def list_type(
-    type: sqltypes.DuckDBPyType, *, connection: DuckDBPyConnection | None = None
-) -> sqltypes.DuckDBPyType: ...
+def list_type(type: IntoDType, *, connection: DuckDBPyConnection | None = None) -> sqltypes.DuckDBPyType: ...
 def load_extension(extension: str, *, connection: DuckDBPyConnection | None = None) -> None: ...
 def map_type(
-    key: sqltypes.DuckDBPyType,
-    value: sqltypes.DuckDBPyType,
-    *,
-    connection: DuckDBPyConnection | None = None,
+    key: IntoDType, value: IntoDType, *, connection: DuckDBPyConnection | None = None
 ) -> sqltypes.DuckDBPyType: ...
 def order(
     df: pandas.DataFrame, order_expr: str, *, connection: DuckDBPyConnection | None = None
@@ -1394,11 +1391,7 @@ def register_filesystem(
 ) -> None: ...
 def remove_function(name: str, *, connection: DuckDBPyConnection | None = None) -> DuckDBPyConnection: ...
 def rollback(*, connection: DuckDBPyConnection | None = None) -> DuckDBPyConnection: ...
-def row_type(
-    fields: dict[str, sqltypes.DuckDBPyType] | lst[sqltypes.DuckDBPyType],
-    *,
-    connection: DuckDBPyConnection | None = None,
-) -> sqltypes.DuckDBPyType: ...
+def row_type(fields: IntoNestedDType, *, connection: DuckDBPyConnection | None = None) -> sqltypes.DuckDBPyType: ...
 def rowcount(*, connection: DuckDBPyConnection | None = None) -> int: ...
 def set_default_connection(connection: DuckDBPyConnection) -> None: ...
 def sql(
@@ -1410,11 +1403,7 @@ def sql(
 ) -> DuckDBPyRelation: ...
 def sqltype(type_str: str, *, connection: DuckDBPyConnection | None = None) -> sqltypes.DuckDBPyType: ...
 def string_type(collation: str = "", *, connection: DuckDBPyConnection | None = None) -> sqltypes.DuckDBPyType: ...
-def struct_type(
-    fields: dict[str, sqltypes.DuckDBPyType] | lst[sqltypes.DuckDBPyType],
-    *,
-    connection: DuckDBPyConnection | None = None,
-) -> sqltypes.DuckDBPyType: ...
+def struct_type(fields: IntoNestedDType, *, connection: DuckDBPyConnection | None = None) -> sqltypes.DuckDBPyType: ...
 def table(table_name: str, *, connection: DuckDBPyConnection | None = None) -> DuckDBPyRelation: ...
 def table_function(
     name: str,
@@ -1426,11 +1415,7 @@ def tf(*, connection: DuckDBPyConnection | None = None) -> dict[str, typing.Any]
 def tokenize(query: str) -> lst[tuple[int, token_type]]: ...
 def torch(*, connection: DuckDBPyConnection | None = None) -> dict[str, typing.Any]: ...
 def type(type_str: str, *, connection: DuckDBPyConnection | None = None) -> sqltypes.DuckDBPyType: ...
-def union_type(
-    members: dict[str, sqltypes.DuckDBPyType] | lst[sqltypes.DuckDBPyType],
-    *,
-    connection: DuckDBPyConnection | None = None,
-) -> sqltypes.DuckDBPyType: ...
+def union_type(members: IntoNestedDType, *, connection: DuckDBPyConnection | None = None) -> sqltypes.DuckDBPyType: ...
 def unregister(view_name: str, *, connection: DuckDBPyConnection | None = None) -> DuckDBPyConnection: ...
 def unregister_filesystem(name: str, *, connection: DuckDBPyConnection | None = None) -> None: ...
 def values(*args: IntoValues, connection: DuckDBPyConnection | None = None) -> DuckDBPyRelation: ...

_duckdb-stubs/_expression.pyi

@@ -1,8 +1,7 @@
 from typing import TYPE_CHECKING, Any, overload
-from duckdb import sqltypes
 
 if TYPE_CHECKING:
-    from ._typing import IntoExpr
+    from ._typing import IntoExpr, IntoDType
 
 class Expression:
     def __add__(self, other: IntoExpr) -> Expression: ...
@@ -40,7 +39,7 @@ class Expression:
     def alias(self, name: str) -> Expression: ...
     def asc(self) -> Expression: ...
     def between(self, lower: IntoExpr, upper: IntoExpr) -> Expression: ...
-    def cast(self, type: sqltypes.DuckDBPyType) -> Expression: ...
+    def cast(self, type: IntoDType) -> Expression: ...
     def collate(self, collation: str) -> Expression: ...
     def desc(self) -> Expression: ...
     def get_name(self) -> str: ...

_duckdb-stubs/_typing.pyi

@@ -1,4 +1,5 @@
 from __future__ import annotations
+
 from typing import TypeAlias, TYPE_CHECKING, Protocol, Any, TypeVar, Generic
 from datetime import date, datetime, time, timedelta
 from decimal import Decimal
@@ -7,6 +8,7 @@ from collections.abc import Mapping, Iterator
 
 if TYPE_CHECKING:
     from ._expression import Expression
+    from ._sqltypes import DuckDBPyType
 
 _T_co = TypeVar("_T_co", covariant=True)
 _S_co = TypeVar("_S_co", bound=tuple[Any, ...], covariant=True)
@@ -54,7 +56,17 @@ PythonLiteral: TypeAlias = (
 # the field_ids argument to to_parquet and write_parquet has a recursive structure
 ParquetFieldIdsType: TypeAlias = Mapping[str, int | ParquetFieldIdsType]
 IntoValues: TypeAlias = list[PythonLiteral] | tuple[Expression, ...] | Expression
-"""Types that can be converted to a table of values."""
+"""Types that can be converted to a table."""
+IntoDType: TypeAlias = DuckDBPyType | str
+"""Types that can be converted to a `DuckDBPyType`.
+
+Passing `INTEGER` is equivalent to passing `DuckDBPyType("INTEGER")` or `DuckDBPyType.INTEGER`.
+
+Note:
+    A `StrEnum` will be handled the same way as a `str`.
+"""
+IntoNestedDType: TypeAlias = dict[str, IntoDType] | list[IntoDType]
+"""Types that can be converted to a nested `DuckDBPyType` (e.g. for struct or union types)."""
 IntoExprColumn: TypeAlias = Expression | str
 """Types that are, or can be used as a `ColumnExpression`."""
 IntoExpr: TypeAlias = IntoExprColumn | PythonLiteral