QA: Migrate type checker from mypy to ty

cratedb_toolkit/__init__.py

@@ -2,7 +2,7 @@
 try:
     from importlib.metadata import PackageNotFoundError, version
 except (ImportError, ModuleNotFoundError):  # pragma:nocover
-    from importlib_metadata import PackageNotFoundError, version  # type: ignore[assignment,no-redef,unused-ignore]
+    from importlib_metadata import PackageNotFoundError, version
 
 __appname__ = "cratedb-toolkit"
 

cratedb_toolkit/adapter/pymongo/api.py

@@ -17,7 +17,7 @@ def __init__(self, dburi: str):
         self.cratedb = DatabaseAdapter(dburi=dburi)
         self.collection_backup = pymongo.collection.Collection
 
-        collection_patched = collection_factory(cratedb=self.cratedb)  # type: ignore[misc]
+        collection_patched = collection_factory(cratedb=self.cratedb)
         self.patches = [
             # Patch PyMongo's `Collection` implementation.
             patch("pymongo.collection.Collection", collection_patched),

cratedb_toolkit/adapter/pymongo/collection.py

@@ -52,7 +52,7 @@ def get_df_info(df: pd.DataFrame) -> str:
         def insert_one(
             self: Collection,
             document: Union[_DocumentType, RawBSONDocument],
-            bypass_document_validation: bool = False,
+            bypass_document_validation: Optional[bool] = None,
             session: Optional[ClientSession] = None,
             comment: Optional[Any] = None,
         ) -> InsertOneResult:
@@ -112,7 +112,7 @@ def insert_many(
             self,
             documents: Iterable[Union[_DocumentType, RawBSONDocument]],
             ordered: bool = True,
-            bypass_document_validation: bool = False,
+            bypass_document_validation: Optional[bool] = None,
             session: Optional[ClientSession] = None,
             comment: Optional[Any] = None,
         ) -> InsertManyResult:
@@ -127,8 +127,8 @@ def gen() -> Iterator[Mapping[str, Any]]:
                     if "_id" in document:
                         identifier = ObjectId(document["_id"])
                         if isinstance(document, RawBSONDocument):
-                            document = document.decode()
-                        del document["_id"]
+                            document = document.decode()  # ty: ignore[unresolved-attribute]
+                        del document["_id"]  # ty: ignore[not-subscriptable]
                     else:
                         identifier = ObjectId()
                     inserted_ids.append(identifier)

cratedb_toolkit/adapter/pymongo/cursor.py

@@ -1,5 +1,4 @@
 # Compansate pymongo<>4.9 woes.
-# mypy: disable-error-code="arg-type,attr-defined,call-arg,misc"
 # Make Python 3.7 and 3.8 support generic types like `dict` instead of `typing.Dict`.
 from __future__ import annotations
 
@@ -250,7 +249,7 @@ def _refresh(self) -> int:
                     self.__query_spec(),
                     self.__projection,
                     self.__codec_options,
-                    self._read_preference(),
+                    self._read_preference(),  # ty: ignore[call-non-callable,invalid-argument-type,missing-argument]
                     self.__limit,
                     self.__batch_size,
                     self.__read_concern,
@@ -275,7 +274,7 @@ def _refresh(self) -> int:
                     limit,
                     self.__id,
                     self.__codec_options,
-                    self._read_preference(),
+                    self._read_preference(),  # ty: ignore[call-non-callable,invalid-argument-type,missing-argument]
                     self.__session,
                     self.__collection.database.client,
                     self.__max_await_time_ms,

cratedb_toolkit/adapter/pymongo/reactor.py

@@ -5,22 +5,23 @@
 from jessiql import Query, QueryObject, QueryObjectDict
 from jessiql.exc import InvalidColumnError
 from jessiql.typing import SARowDict
+from sqlalchemy.orm import Mapper, registry
 
 
-def table_to_model(table: sa.Table) -> t.Type[sa.orm.Mapper]:
+def table_to_model(table: sa.Table) -> t.Type[Mapper]:
     """
     Create SQLAlchemy model class from Table object.
 
     - https://docs.sqlalchemy.org/en/14/orm/mapping_styles.html#imperative-mapping
     - https://sparrigan.github.io/sql/sqla/2016/01/03/dynamic-tables.html
     """
-    mapper_registry = sa.orm.registry(metadata=table.metadata)
+    mapper_registry = registry(metadata=table.metadata)
     Surrogate = type("Surrogate", (), {})
     mapper_registry.map_imperatively(Surrogate, table)
-    return Surrogate
+    return Surrogate  # ty: ignore[invalid-return-type]
 
 
-def reflect_model(engine: t.Any, metadata: sa.MetaData, table_name: str) -> t.Type[sa.orm.Mapper]:
+def reflect_model(engine: t.Any, metadata: sa.MetaData, table_name: str) -> t.Type[Mapper]:
     """
     Create SQLAlchemy model class by reflecting a database table.
     """
@@ -29,7 +30,7 @@ def reflect_model(engine: t.Any, metadata: sa.MetaData, table_name: str) -> t.Ty
 
 
 def mongodb_query(
-    model: t.Type[sa.orm.Mapper],
+    model: t.Type[Mapper],
     select: t.Union[t.List, None] = None,
     filter: t.Union[t.Dict[str, t.Any], None] = None,  # noqa: A002
     sort: t.Union[t.List[str], None] = None,
@@ -38,7 +39,7 @@ def mongodb_query(
     Create a JessiQL Query object from an SQLAlchemy model class and typical MongoDB query parameters.
     """
 
-    select = select or list(model._sa_class_manager.keys())  # type: ignore[attr-defined]
+    select = select or list(model._sa_class_manager.keys())
 
     filter = filter or {}  # noqa: A001
     sort = sort or []

cratedb_toolkit/adapter/pymongo/util.py

@@ -26,7 +26,7 @@ def from_str(cls, oid: str):
         return cls(bytes(oid, "ascii"))
 
     def __str__(self) -> str:
-        return binascii.hexlify(self.__id).decode()  # type: ignore[arg-type]
+        return binascii.hexlify(self.__id).decode()  # ty: ignore[invalid-argument-type]
 
     def __repr__(self) -> str:
         return f"ObjectId('{self!s}')"

cratedb_toolkit/adapter/rockset/cli.py

@@ -25,7 +25,7 @@ def help_serve():
     """  # noqa: E501
 
 
-@click.group(cls=ClickAliasedGroup)  # type: ignore[arg-type]
+@click.group(cls=ClickAliasedGroup)
 @option_cluster_url
 @click.option("--verbose", is_flag=True, required=False, help="Turn on logging")
 @click.option("--debug", is_flag=True, required=False, help="Turn on logging with debug level")

cratedb_toolkit/adapter/rockset/server/api/query.py

@@ -52,7 +52,7 @@ async def execute(request: Request, adapter: t.Annotated[DatabaseAdapter, Depend
 
     results = adapter.run_sql(
         sql=sql,
-        parameters=parameters,
+        parameters=parameters or {},
         records=True,
     )
     time_duration = time.time_ns() - time_start

cratedb_toolkit/cfr/cli.py

@@ -85,7 +85,7 @@ def help_statistics():
     """  # noqa: E501
 
 
-@click.group(cls=ClickAliasedGroup, help=docstring_format_verbatim(help_statistics.__doc__))  # type: ignore[arg-type]
+@click.group(cls=ClickAliasedGroup, help=docstring_format_verbatim(help_statistics.__doc__))
 @click.pass_context
 def job_statistics(ctx: click.Context):
     """
@@ -221,7 +221,7 @@ def job_statistics_ui(ctx: click.Context):
     uvicorn.run(app, host="localhost", port=7777, log_level="info")
 
 
-@click.group(cls=ClickAliasedGroup)  # type: ignore[arg-type]
+@click.group(cls=ClickAliasedGroup)
 @click.pass_context
 def info(ctx: click.Context):
     """

cratedb_toolkit/cfr/marimo.py

@@ -1,5 +1,4 @@
 # ruff: noqa: B018, ERA001, S608, T201
-# mypy: disable-error-code="arg-type"
 
 # Source:
 # https://docs.marimo.io/guides/working_with_data/dataframes.html

cratedb_toolkit/cfr/systable.py

@@ -83,7 +83,9 @@ def __init__(self, dburi: str):
     def table_names(self):
         return self.inspector.get_table_names(schema=SystemTableKnowledge.SYS_SCHEMA)
 
-    def ddl(self, tablename_in: str, tablename_out: str, out_schema: str = None, with_drop_table: bool = False) -> str:
+    def ddl(
+        self, tablename_in: str, tablename_out: str, out_schema: t.Optional[str] = None, with_drop_table: bool = False
+    ) -> str:
         meta = sa.MetaData(schema=SystemTableKnowledge.SYS_SCHEMA)
         table = sa.Table(tablename_in, meta, autoload_with=self.engine)
         table.schema = out_schema
@@ -147,9 +149,9 @@ def dump_table(self, frame: "pl.DataFrame", file: t.Union[t.TextIO, None] = None
             # return df.write_csv()  # noqa: ERA001
             return frame.to_pandas().to_csv(file)
         elif self.data_format in ["jsonl", "ndjson"]:
-            return frame.write_ndjson(file and file.buffer)  # type: ignore[arg-type]
+            return frame.write_ndjson(file and file.buffer)
         elif self.data_format in ["parquet", "pq"]:
-            return frame.write_parquet(file and file.buffer)  # type: ignore[arg-type]
+            return frame.write_parquet(file and file.buffer)  # ty: ignore[invalid-argument-type]
         else:
             raise NotImplementedError(f"Output format not implemented: {self.data_format}")
 

cratedb_toolkit/cli.py

@@ -18,7 +18,7 @@
 from .util.setting import init_dotenv
 
 
-@click.group(cls=ClickAliasedGroup)  # type: ignore[arg-type]
+@click.group(cls=ClickAliasedGroup)
 @click.option("--verbose", is_flag=True, required=False, help="Turn on logging")
 @click.option("--debug", is_flag=True, required=False, help="Turn on logging with debug level")
 @click.version_option()

cratedb_toolkit/cluster/cli.py

@@ -57,7 +57,7 @@ def help_list_jobs():
     """
 
 
-@click.group(cls=ClickAliasedGroup)  # type: ignore[arg-type]
+@click.group(cls=ClickAliasedGroup)
 @click.option("--verbose", is_flag=True, required=False, help="Turn on logging")
 @click.option("--debug", is_flag=True, required=False, help="Turn on logging with debug level")
 @click.version_option()

cratedb_toolkit/cluster/core.py

@@ -9,6 +9,7 @@
 from pathlib import Path
 
 import click
+from crate.client.connection import Connection
 
 from cratedb_toolkit.cluster.croud import CloudClusterServices, CloudRootServices
 from cratedb_toolkit.cluster.guide import DataImportGuide
@@ -110,11 +111,11 @@ class ManagedCluster(ClusterBase):
 
     def __init__(
         self,
-        cluster_id: str = None,
-        cluster_name: str = None,
-        settings: ManagedClusterSettings = None,
-        address: DatabaseAddress = None,
-        info: ClusterInformation = None,
+        cluster_id: t.Optional[str] = None,
+        cluster_name: t.Optional[str] = None,
+        settings: t.Optional[ManagedClusterSettings] = None,
+        address: t.Optional[DatabaseAddress] = None,
+        info: t.Optional[ClusterInformation] = None,
         stop_on_exit: bool = False,
     ):
         super().__init__()
@@ -408,7 +409,7 @@ def adapter(self) -> DatabaseAdapter:
             raise DatabaseAddressMissingError()
         return DatabaseAdapter(dburi=self.address.dburi, jwt=self.info.jwt)
 
-    def get_client_bundle(self, username: str = None, password: str = None) -> ClientBundle:
+    def get_client_bundle(self, username: t.Optional[str] = None, password: t.Optional[str] = None) -> ClientBundle:
         """
         Return a bundle of client handles to the CrateDB Cloud cluster database.
 
@@ -431,7 +432,7 @@ def get_client_bundle(self, username: str = None, password: str = None) -> Clien
         adapter = DatabaseAdapter(address.dburi)
         self._client_bundle = ClientBundle(
             adapter=adapter,
-            dbapi=adapter.connection.connection.dbapi_connection,
+            dbapi=t.cast(Connection, adapter.connection.connection.dbapi_connection),
             sqlalchemy=adapter.engine,
         )
         return self._client_bundle
@@ -492,7 +493,7 @@ def adapter(self) -> DatabaseAdapter:
         """
         return DatabaseAdapter(dburi=self.address.dburi)
 
-    def get_client_bundle(self, username: str = None, password: str = None) -> ClientBundle:
+    def get_client_bundle(self, username: t.Optional[str] = None, password: t.Optional[str] = None) -> ClientBundle:
         """
         Return a bundle of client handles to the CrateDB Cloud cluster database.
 
@@ -512,7 +513,7 @@ def get_client_bundle(self, username: str = None, password: str = None) -> Clien
         adapter = DatabaseAdapter(address.dburi)
         self._client_bundle = ClientBundle(
             adapter=adapter,
-            dbapi=adapter.connection.connection.dbapi_connection,
+            dbapi=t.cast(Connection, adapter.connection.connection.dbapi_connection),
             sqlalchemy=adapter.engine,
         )
         return self._client_bundle
@@ -592,7 +593,10 @@ def from_options(cls, options: ClusterAddressOptions) -> t.Union[ManagedCluster,
 
     @classmethod
     def create(
-        cls, cluster_id: str = None, cluster_name: str = None, cluster_url: str = None
+        cls,
+        cluster_id: t.Optional[str] = None,
+        cluster_name: t.Optional[str] = None,
+        cluster_url: t.Optional[str] = None,
     ) -> t.Union[ManagedCluster, StandaloneCluster]:
         """
         Create the cluster instance based on the provided parameters.

cratedb_toolkit/cluster/croud.py

@@ -164,7 +164,7 @@ def list_projects(self):
             item["backup_location"] = _transform_backup_location(item["backup_location"])
         return data
 
-    def create_project(self, name: str, organization_id: str = None):
+    def create_project(self, name: str, organization_id: t.Optional[str] = None):
         """
         Create project.
 
@@ -184,7 +184,7 @@ def create_project(self, name: str, organization_id: str = None):
 
         call = CroudCall(
             fun=project_create,
-            specs=command_tree["projects"]["commands"]["create"]["extra_args"],
+            specs=command_tree["projects"]["commands"]["create"]["extra_args"],  # ty: ignore[invalid-argument-type,not-subscriptable]
             arguments=arguments,
         )
 
@@ -216,7 +216,7 @@ def get_or_create_project(self, name: str) -> str:
 
         return project_id
 
-    def deploy_cluster(self, name: str, project_id: str, subscription_id: str = None):
+    def deploy_cluster(self, name: str, project_id: str, subscription_id: t.Optional[str] = None):
         """
         Deploy cluster.
 
@@ -273,7 +273,7 @@ def deploy_cluster(self, name: str, project_id: str, subscription_id: str = None
 
         call = CroudCall(
             fun=clusters_deploy,
-            specs=command_tree["clusters"]["commands"]["deploy"]["extra_args"],
+            specs=command_tree["clusters"]["commands"]["deploy"]["extra_args"],  # ty: ignore[invalid-argument-type,not-subscriptable]
             arguments=[
                 "--subscription-id",
                 subscription_id,

cratedb_toolkit/cmd/tail/main.py

@@ -5,10 +5,10 @@
 import typing as t
 
 import attr
-import colorlog
 import orjson
 import sqlparse
 import yaml
+from colorlog.escape_codes import escape_codes
 
 from cratedb_toolkit.model import TableAddress
 from cratedb_toolkit.util.database import DatabaseAdapter
@@ -37,9 +37,9 @@ def template(self) -> str:
 
     @property
     def label(self):
-        red = colorlog.escape_codes.escape_codes["red"]
-        green = colorlog.escape_codes.escape_codes["green"]
-        reset = colorlog.escape_codes.escape_codes["reset"]
+        red = escape_codes["red"]
+        green = escape_codes["green"]
+        reset = escape_codes["reset"]
         if self.error:
             return f"{red}ERROR{reset}"
         else:

cratedb_toolkit/datasets/model.py

@@ -9,7 +9,7 @@
 try:
     from typing import Literal
 except ImportError:
-    from typing_extensions import Literal  # type: ignore[assignment]
+    from typing_extensions import Literal
 
 
 @dataclasses.dataclass

cratedb_toolkit/datasets/util.py

@@ -1,9 +1,10 @@
 import os
+import typing as t
 import zipfile
 from unittest.mock import patch
 
 
-def load_dataset_kaggle(dataset: str, path: str, file_name: str = None):
+def load_dataset_kaggle(dataset: str, path: str, file_name: t.Optional[str] = None):
     """
     Download complete dataset or individual files from Kaggle.
 

cratedb_toolkit/dms/cli.py

@@ -17,7 +17,7 @@ def help_table_mappings():
     """
 
 
-@click.group(cls=ClickAliasedGroup, help="AWS DMS utilities")  # type: ignore[arg-type]
+@click.group(cls=ClickAliasedGroup, help="AWS DMS utilities")
 @click.pass_context
 def cli(ctx: click.Context):
     pass

cratedb_toolkit/dms/table_mapping.py

@@ -42,8 +42,8 @@ def add_rule(
         action: str,
         locator: Dict[str, str],
         id: Union[str, None] = None,  # noqa: A002
-        filters: List[Any] = None,
-        mapping_parameters: Dict[str, str] = None,
+        filters: Union[List[Any], None] = None,
+        mapping_parameters: Union[Dict[str, str], None] = None,
     ) -> "TableMappingBuilder":
         if id is None:
             id = str(len(self.rules) + 1)  # noqa: A001