Ruff ANN401: specialized Any type hints where possible

Author: evertlammerts

duckdb/bytes_io_wrapper.py

@@ -48,7 +48,7 @@ def __init__(self, buffer: Union[StringIO, TextIOBase], encoding: str = "utf-8")
         # overflow to the front of the bytestring the next time reading is performed
         self.overflow = b""
 
-    def __getattr__(self, attr: str) -> Any:  # noqa: D105
+    def __getattr__(self, attr: str) -> Any:  # noqa: D105, ANN401
         return getattr(self.buffer, attr)
 
     def read(self, n: Union[int, None] = -1) -> bytes:  # noqa: D102

duckdb/experimental/spark/sql/column.py

@@ -1,4 +1,5 @@
-from typing import TYPE_CHECKING, Any, Callable, Union, cast  # noqa: D100
+from collections.abc import Iterable  # noqa: D100
+from typing import TYPE_CHECKING, Any, Callable, Union, cast
 
 from ..exception import ContributionsAcceptedError
 from .types import DataType
@@ -136,7 +137,7 @@ def __neg__(self) -> "Column":  # noqa: D105
 
     __rpow__ = _bin_op("__rpow__")
 
-    def __getitem__(self, k: Any) -> "Column":
+    def __getitem__(self, k: Any) -> "Column":  # noqa: ANN401
         """An expression that gets an item at position ``ordinal`` out of a list,
         or gets an item by key out of a dict.
 
@@ -176,7 +177,7 @@ def __getitem__(self, k: Any) -> "Column":
             expr_str = str(self.expr) + "." + str(k)
             return Column(ColumnExpression(expr_str))
 
-    def __getattr__(self, item: Any) -> "Column":
+    def __getattr__(self, item: Any) -> "Column":  # noqa: ANN401
         """An expression that gets an item at position ``ordinal`` out of a list,
         or gets an item by key out of a dict.
 
@@ -208,15 +209,15 @@ def __getattr__(self, item: Any) -> "Column":
     def alias(self, alias: str) -> "Column":  # noqa: D102
         return Column(self.expr.alias(alias))
 
-    def when(self, condition: "Column", value: Any) -> "Column":  # noqa: D102
+    def when(self, condition: "Column", value: Union["Column", str]) -> "Column":  # noqa: D102
         if not isinstance(condition, Column):
             msg = "condition should be a Column"
             raise TypeError(msg)
         v = _get_expr(value)
         expr = self.expr.when(condition.expr, v)
         return Column(expr)
 
-    def otherwise(self, value: Any) -> "Column":  # noqa: D102
+    def otherwise(self, value: Union["Column", str]) -> "Column":  # noqa: D102
         v = _get_expr(value)
         expr = self.expr.otherwise(v)
         return Column(expr)
@@ -229,7 +230,7 @@ def cast(self, dataType: Union[DataType, str]) -> "Column":  # noqa: D102
             internal_type = dataType.duckdb_type
         return Column(self.expr.cast(internal_type))
 
-    def isin(self, *cols: Any) -> "Column":  # noqa: D102
+    def isin(self, *cols: Union[Iterable[Union["Column", str]], Union["Column", str]]) -> "Column":  # noqa: D102
         if len(cols) == 1 and isinstance(cols[0], (list, set)):
             # Only one argument supplied, it's a list
             cols = cast("tuple", cols[0])

duckdb/experimental/spark/sql/dataframe.py

@@ -282,7 +282,7 @@ def withColumnsRenamed(self, colsMap: dict[str, str]) -> "DataFrame":
         rel = self.relation.select(*cols)
         return DataFrame(rel, self.session)
 
-    def transform(self, func: Callable[..., "DataFrame"], *args: Any, **kwargs: Any) -> "DataFrame":
+    def transform(self, func: Callable[..., "DataFrame"], *args: Any, **kwargs: Any) -> "DataFrame":  # noqa: ANN401
         """Returns a new :class:`DataFrame`. Concise syntax for chaining custom transformations.
 
         .. versionadded:: 3.0.0
@@ -342,7 +342,7 @@ def transform(self, func: Callable[..., "DataFrame"], *args: Any, **kwargs: Any)
         )
         return result
 
-    def sort(self, *cols: Union[str, Column, list[Union[str, Column]]], **kwargs: Any) -> "DataFrame":
+    def sort(self, *cols: Union[str, Column, list[Union[str, Column]]], **kwargs: Any) -> "DataFrame":  # noqa: ANN401
         """Returns a new :class:`DataFrame` sorted by the specified column(s).
 
         Parameters

duckdb/experimental/spark/sql/functions.py

@@ -90,7 +90,7 @@ def ucase(str: "ColumnOrName") -> Column:
     return upper(str)
 
 
-def when(condition: "Column", value: Any) -> Column:  # noqa: D103
+def when(condition: "Column", value: Column | str) -> Column:  # noqa: D103
     if not isinstance(condition, Column):
         msg = "condition should be a Column"
         raise TypeError(msg)
@@ -143,7 +143,7 @@ def array(*cols: Union["ColumnOrName", Union[list["ColumnOrName"], tuple["Column
     return _invoke_function_over_columns("list_value", *cols)
 
 
-def lit(col: Any) -> Column:  # noqa: D103
+def lit(col: Any) -> Column:  # noqa: D103, ANN401
     return col if isinstance(col, Column) else Column(ConstantExpression(col))
 
 
@@ -842,7 +842,7 @@ def collect_list(col: "ColumnOrName") -> Column:
     return array_agg(col)
 
 
-def array_append(col: "ColumnOrName", value: Any) -> Column:
+def array_append(col: "ColumnOrName", value: Column | str) -> Column:
     """Collection function: returns an array of the elements in col1 along
     with the added element in col2 at the last of the array.
 
@@ -876,7 +876,7 @@ def array_append(col: "ColumnOrName", value: Any) -> Column:
     return _invoke_function("list_append", _to_column_expr(col), _get_expr(value))
 
 
-def array_insert(arr: "ColumnOrName", pos: Union["ColumnOrName", int], value: Any) -> Column:
+def array_insert(arr: "ColumnOrName", pos: Union["ColumnOrName", int], value: Column | str) -> Column:
     """Collection function: adds an item into a given array at a specified array index.
     Array indices start at 1, or start from the end if index is negative.
     Index above array size appends the array, or prepends the array if index is negative,
@@ -969,7 +969,7 @@ def array_insert(arr: "ColumnOrName", pos: Union["ColumnOrName", int], value: An
     )
 
 
-def array_contains(col: "ColumnOrName", value: Any) -> Column:
+def array_contains(col: "ColumnOrName", value: Column | str) -> Column:
     """Collection function: returns null if the array is null, true if the array contains the
     given value, and false otherwise.
 
@@ -1937,7 +1937,7 @@ def array_compact(col: "ColumnOrName") -> Column:
     )
 
 
-def array_remove(col: "ColumnOrName", element: Any) -> Column:
+def array_remove(col: "ColumnOrName", element: Any) -> Column:  # noqa: ANN401
     """Collection function: Remove all elements that equal to element from the given array.
 
     .. versionadded:: 2.4.0
@@ -5083,7 +5083,7 @@ def array_join(col: "ColumnOrName", delimiter: str, null_replacement: Optional[s
     return _invoke_function("array_to_string", col, ConstantExpression(delimiter))
 
 
-def array_position(col: "ColumnOrName", value: Any) -> Column:
+def array_position(col: "ColumnOrName", value: Any) -> Column:  # noqa: ANN401
     """Collection function: Locates the position of the first occurrence of the given value
     in the given array. Returns null if either of the arguments are null.
 
@@ -5122,7 +5122,7 @@ def array_position(col: "ColumnOrName", value: Any) -> Column:
     )
 
 
-def array_prepend(col: "ColumnOrName", value: Any) -> Column:
+def array_prepend(col: "ColumnOrName", value: Any) -> Column:  # noqa: ANN401
     """Collection function: Returns an array containing element as
     well as all elements from array. The new element is positioned
     at the beginning of the array.

duckdb/experimental/spark/sql/session.py

@@ -1,6 +1,6 @@
 import uuid  # noqa: D100
 from collections.abc import Iterable, Sized
-from typing import TYPE_CHECKING, Any, Optional, Union
+from typing import TYPE_CHECKING, Any, NoReturn, Optional, Union
 
 if TYPE_CHECKING:
     from pandas.core.frame import DataFrame as PandasDataFrame
@@ -205,7 +205,7 @@ def range(  # noqa: D102
 
         return DataFrame(self.conn.table_function("range", parameters=[start, end, step]), self)
 
-    def sql(self, sqlQuery: str, **kwargs: Any) -> DataFrame:  # noqa: D102
+    def sql(self, sqlQuery: str, **kwargs: Any) -> DataFrame:  # noqa: D102, ANN401
         if kwargs:
             raise NotImplementedError
         relation = self.conn.sql(sqlQuery)
@@ -246,7 +246,7 @@ def sparkContext(self) -> SparkContext:  # noqa: D102
         return self._context
 
     @property
-    def streams(self) -> Any:  # noqa: D102
+    def streams(self) -> NoReturn:  # noqa: D102
         raise ContributionsAcceptedError
 
     @property
@@ -278,7 +278,10 @@ def getOrCreate(self) -> "SparkSession":  # noqa: D102
             return SparkSession(context)
 
         def config(  # noqa: D102
-            self, key: Optional[str] = None, value: Optional[Any] = None, conf: Optional[SparkConf] = None
+            self,
+            key: Optional[str] = None,
+            value: Optional[Any] = None,  # noqa: ANN401
+            conf: Optional[SparkConf] = None,
         ) -> "SparkSession.Builder":
             return self
 

duckdb/experimental/spark/sql/types.py

@@ -1,4 +1,5 @@
-# This code is based on code from Apache Spark under the license found in the LICENSE  # noqa: D100
+# ruff: noqa: D100
+# This code is based on code from Apache Spark under the license found in the LICENSE
 # file located in the 'spark' folder.
 
 import calendar
@@ -11,6 +12,7 @@
 from typing import (
     Any,
     ClassVar,
+    NoReturn,
     Optional,
     TypeVar,
     Union,
@@ -102,11 +104,11 @@ def needConversion(self) -> bool:
         """
         return False
 
-    def toInternal(self, obj: Any) -> Any:
+    def toInternal(self, obj: Any) -> Any:  # noqa: ANN401
         """Converts a Python object into an internal SQL object."""
         return obj
 
-    def fromInternal(self, obj: Any) -> Any:
+    def fromInternal(self, obj: Any) -> Any:  # noqa: ANN401
         """Converts an internal SQL object into a native Python object."""
         return obj
 
@@ -889,7 +891,7 @@ def simpleString(self) -> str:  # noqa: D102
     def __repr__(self) -> str:  # noqa: D105
         return "StructType([%s])" % ", ".join(str(field) for field in self)
 
-    def __contains__(self, item: Any) -> bool:  # noqa: D105
+    def __contains__(self, item: str) -> bool:  # noqa: D105
         return item in self.names
 
     def extract_types_and_names(self) -> tuple[list[str], list[str]]:  # noqa: D102
@@ -1010,21 +1012,21 @@ def _cachedSqlType(cls) -> DataType:
             cls._cached_sql_type = cls.sqlType()  # type: ignore[attr-defined]
         return cls._cached_sql_type  # type: ignore[attr-defined]
 
-    def toInternal(self, obj: Any) -> Any:
+    def toInternal(self, obj: Any) -> Any:  # noqa: ANN401
         if obj is not None:
             return self._cachedSqlType().toInternal(self.serialize(obj))
 
-    def fromInternal(self, obj: Any) -> Any:
+    def fromInternal(self, obj: Any) -> Any:  # noqa: ANN401
         v = self._cachedSqlType().fromInternal(obj)
         if v is not None:
             return self.deserialize(v)
 
-    def serialize(self, obj: Any) -> Any:
+    def serialize(self, obj: Any) -> NoReturn:  # noqa: ANN401
         """Converts a user-type object into a SQL datum."""
         msg = "UDT must implement toInternal()."
         raise NotImplementedError(msg)
 
-    def deserialize(self, datum: Any) -> Any:
+    def deserialize(self, datum: Any) -> NoReturn:  # noqa: ANN401
         """Converts a SQL datum into a user-type object."""
         msg = "UDT must implement fromInternal()."
         raise NotImplementedError(msg)
@@ -1132,7 +1134,7 @@ class Row(tuple):
     def __new__(cls, *args: str) -> "Row": ...
 
     @overload
-    def __new__(cls, **kwargs: Any) -> "Row": ...
+    def __new__(cls, **kwargs: Any) -> "Row": ...  # noqa: ANN401
 
     def __new__(cls, *args: Optional[str], **kwargs: Optional[Any]) -> "Row":  # noqa: D102
         if args and kwargs:
@@ -1179,7 +1181,7 @@ def asDict(self, recursive: bool = False) -> dict[str, Any]:
 
         if recursive:
 
-            def conv(obj: Any) -> Any:
+            def conv(obj: Row | list | dict | object) -> list | dict | object:
                 if isinstance(obj, Row):
                     return obj.asDict(True)
                 elif isinstance(obj, list):
@@ -1193,22 +1195,22 @@ def conv(obj: Any) -> Any:
         else:
             return dict(zip(self.__fields__, self))
 
-    def __contains__(self, item: Any) -> bool:  # noqa: D105
+    def __contains__(self, item: Any) -> bool:  # noqa: D105, ANN401
         if hasattr(self, "__fields__"):
             return item in self.__fields__
         else:
             return super(Row, self).__contains__(item)
 
     # let object acts like class
-    def __call__(self, *args: Any) -> "Row":
+    def __call__(self, *args: Any) -> "Row":  # noqa: ANN401
         """Create new Row object."""
         if len(args) > len(self):
             raise ValueError(
                 "Can not create Row with fields %s, expected %d values but got %s" % (self, len(self), args)
             )
         return _create_row(self, args)
 
-    def __getitem__(self, item: Any) -> Any:  # noqa: D105
+    def __getitem__(self, item: Any) -> Any:  # noqa: D105, ANN401
         if isinstance(item, (int, slice)):
             return super(Row, self).__getitem__(item)
         try:
@@ -1221,7 +1223,7 @@ def __getitem__(self, item: Any) -> Any:  # noqa: D105
         except ValueError:
             raise ValueError(item)
 
-    def __getattr__(self, item: str) -> Any:  # noqa: D105
+    def __getattr__(self, item: str) -> Any:  # noqa: D105, ANN401
         if item.startswith("__"):
             raise AttributeError(item)
         try:
@@ -1234,7 +1236,7 @@ def __getattr__(self, item: str) -> Any:  # noqa: D105
         except ValueError:
             raise AttributeError(item)
 
-    def __setattr__(self, key: Any, value: Any) -> None:  # noqa: D105
+    def __setattr__(self, key: Any, value: Any) -> None:  # noqa: D105, ANN401
         if key != "__fields__":
             msg = "Row is read-only"
             raise RuntimeError(msg)