Ruff ANN001: fixed missing annotations for function args

Author: evertlammerts

duckdb/experimental/spark/exception.py

@@ -4,7 +4,7 @@ class ContributionsAcceptedError(NotImplementedError):  # noqa: D100
     feel free to open up a PR or a Discussion over on https://github.com/duckdb/duckdb.
     """  # noqa: D205
 
-    def __init__(self, message=None) -> None:  # noqa: D107
+    def __init__(self, message: str=None) -> None:  # noqa: D107
         doc = self.__class__.__doc__
         if message:
             doc = message + "\n" + doc

duckdb/experimental/spark/sql/catalog.py

@@ -40,7 +40,7 @@ def __init__(self, session: SparkSession) -> None:  # noqa: D107
     def listDatabases(self) -> list[Database]:  # noqa: D102
         res = self._session.conn.sql("select database_name from duckdb_databases()").fetchall()
 
-        def transform_to_database(x) -> Database:
+        def transform_to_database(x: list[str]) -> Database:
             return Database(name=x[0], description=None, locationUri="")
 
         databases = [transform_to_database(x) for x in res]
@@ -49,7 +49,7 @@ def transform_to_database(x) -> Database:
     def listTables(self) -> list[Table]:  # noqa: D102
         res = self._session.conn.sql("select table_name, database_name, sql, temporary from duckdb_tables()").fetchall()
 
-        def transform_to_table(x) -> Table:
+        def transform_to_table(x: list[str]) -> Table:
             return Table(name=x[0], database=x[1], description=x[2], tableType="", isTemporary=x[3])
 
         tables = [transform_to_table(x) for x in res]
@@ -63,7 +63,7 @@ def listColumns(self, tableName: str, dbName: Optional[str] = None) -> list[Colu
             query += f" and database_name = '{dbName}'"
         res = self._session.conn.sql(query).fetchall()
 
-        def transform_to_column(x) -> Column:
+        def transform_to_column(x: list[str|bool]) -> Column:
             return Column(name=x[0], description=None, dataType=x[1], nullable=x[2], isPartition=False, isBucket=False)
 
         columns = [transform_to_column(x) for x in res]

duckdb/experimental/spark/sql/column.py

@@ -12,7 +12,7 @@
 __all__ = ["Column"]
 
 
-def _get_expr(x) -> Expression:
+def _get_expr(x: 'Column' | str) -> Expression:
     return x.expr if isinstance(x, Column) else ConstantExpression(x)
 
 

duckdb/experimental/spark/sql/dataframe.py

@@ -698,7 +698,7 @@ def join(
                 on = str(on)
             assert isinstance(how, str), "how should be a string"
 
-            def map_to_recognized_jointype(how):
+            def map_to_recognized_jointype(how: str):
                 known_aliases = {
                     "inner": [],
                     "outer": ["full", "fullouter", "full_outer"],
@@ -1354,7 +1354,7 @@ def collect(self) -> list[Row]:  # noqa: D102
         columns = self.relation.columns
         result = self.relation.fetchall()
 
-        def construct_row(values, names) -> Row:
+        def construct_row(values: list, names: list[str]) -> Row:
             row = tuple.__new__(Row, list(values))
             row.__fields__ = list(names)
             return row

duckdb/experimental/spark/sql/functions.py

@@ -99,7 +99,7 @@ def when(condition: "Column", value: Any) -> Column:  # noqa: D103
     return Column(expr)
 
 
-def _inner_expr_or_val(val):
+def _inner_expr_or_val(val: Column | str) -> Column | str:
     return val.expr if isinstance(val, Column) else val
 
 

duckdb/experimental/spark/sql/session.py

@@ -1,6 +1,6 @@
 import uuid  # noqa: D100
 from collections.abc import Iterable
-from typing import TYPE_CHECKING, Any, Optional, Union
+from typing import TYPE_CHECKING, Any, Optional, Union, Sized
 
 if TYPE_CHECKING:
     from pandas.core.frame import DataFrame as PandasDataFrame
@@ -58,7 +58,7 @@ def _create_dataframe(self, data: Union[Iterable[Any], "PandasDataFrame"]) -> Da
             self.conn.register(unique_name, data)
             return DataFrame(self.conn.sql(f'select * from "{unique_name}"'), self)
 
-        def verify_tuple_integrity(tuples):
+        def verify_tuple_integrity(tuples: list[tuple]):
             if len(tuples) <= 1:
                 return
             expected_length = len(tuples[0])
@@ -80,8 +80,8 @@ def verify_tuple_integrity(tuples):
             data = list(data)
         verify_tuple_integrity(data)
 
-        def construct_query(tuples) -> str:
-            def construct_values_list(row, start_param_idx):
+        def construct_query(tuples: Iterable) -> str:
+            def construct_values_list(row: Sized, start_param_idx: int):
                 parameter_count = len(row)
                 parameters = [f"${x + start_param_idx}" for x in range(parameter_count)]
                 parameters = "(" + ", ".join(parameters) + ")"
@@ -98,7 +98,7 @@ def construct_values_list(row, start_param_idx):
 
         query = construct_query(data)
 
-        def construct_parameters(tuples):
+        def construct_parameters(tuples: Iterable):
             parameters = []
             for row in tuples:
                 parameters.extend(list(row))
@@ -109,7 +109,7 @@ def construct_parameters(tuples):
         rel = self.conn.sql(query, params=parameters)
         return DataFrame(rel, self)
 
-    def _createDataFrameFromPandas(self, data: "PandasDataFrame", types, names) -> DataFrame:
+    def _createDataFrameFromPandas(self, data: "PandasDataFrame", types: list[str] | None, names: list[str] | None) -> DataFrame:
         df = self._create_dataframe(data)
 
         # Cast to types

duckdb/experimental/spark/sql/types.py

@@ -66,7 +66,7 @@
 class DataType:
     """Base class for data types."""
 
-    def __init__(self, duckdb_type) -> None:  # noqa: D107
+    def __init__(self, duckdb_type: DuckDBPyType) -> None:  # noqa: D107
         self.duckdb_type = duckdb_type
 
     def __repr__(self) -> str:  # noqa: D105

duckdb/filesystem.py

@@ -1,12 +1,13 @@
 from io import TextIOBase  # noqa: D100
+from typing import IO
 
 from fsspec import AbstractFileSystem
 from fsspec.implementations.memory import MemoryFile, MemoryFileSystem
 
 from .bytes_io_wrapper import BytesIOWrapper
 
 
-def is_file_like(obj):  # noqa: D103
+def is_file_like(obj):  # noqa: D103, ANN001
     # We only care that we can read from the file
     return hasattr(obj, "read") and hasattr(obj, "seek")
 
@@ -16,7 +17,7 @@ class ModifiedMemoryFileSystem(MemoryFileSystem):  # noqa: D101
     # defer to the original implementation that doesn't hardcode the protocol
     _strip_protocol = classmethod(AbstractFileSystem._strip_protocol.__func__)
 
-    def add_file(self, object, path):  # noqa: D102
+    def add_file(self, object: IO, path: str):  # noqa: D102
         if not is_file_like(object):
             msg = "Can not read from a non file-like object"
             raise ValueError(msg)

duckdb/udf.py

@@ -1,4 +1,7 @@
-def vectorized(func):  # noqa: D100
+from typing import Callable
+
+
+def vectorized(func: Callable):  # noqa: D100
     """Decorate a function with annotated function parameters, so DuckDB can infer that the function should be provided with pyarrow arrays and should expect pyarrow array(s) as output."""
     import types
     from inspect import signature

duckdb_packaging/_versioning.py

@@ -147,7 +147,9 @@ def strip_post_from_version(version: str) -> str:
     return re.sub(r"[\.-]post[0-9]+", "", version)
 
 
-def get_git_describe(repo_path: Optional[pathlib.Path] = None, since_major=False, since_minor=False) -> Optional[str]:
+def get_git_describe(
+    repo_path: Optional[pathlib.Path] = None, since_major: bool = False, since_minor: bool = False
+) -> Optional[str]:
     """Get git describe output for version determination.
 
     Returns: