Type-check entire code base with mypy

Lots of `# type: ignore` in the internals, but the types should
be decent enough for consumption.

Author: Photonios

.circleci/config.yml

@@ -119,7 +119,7 @@ jobs:
     steps:
       - checkout
       - install-dependencies:
-          extra: analysis
+          extra: analysis, test
       - run:
           name: Verify
           command: python setup.py verify

psqlextra/backend/base.py

@@ -1,5 +1,7 @@
 import logging
 
+from typing import TYPE_CHECKING
+
 from django.conf import settings
 from django.db import ProgrammingError
 
@@ -8,17 +10,31 @@
 from .operations import PostgresOperations
 from .schema import PostgresSchemaEditor
 
+from django.db.backends.postgresql.base import (  # isort:skip
+    DatabaseWrapper as PostgresDatabaseWrapper,
+)
+
+
 logger = logging.getLogger(__name__)
 
 
-class DatabaseWrapper(base_impl.backend()):
+if TYPE_CHECKING:
+
+    class Wrapper(PostgresDatabaseWrapper):
+        pass
+
+else:
+    Wrapper = base_impl.backend()
+
+
+class DatabaseWrapper(Wrapper):
     """Wraps the standard PostgreSQL database back-end.
 
     Overrides the schema editor with our custom schema editor and makes
     sure the `hstore` extension is enabled.
     """
 
-    SchemaEditorClass = PostgresSchemaEditor
+    SchemaEditorClass = PostgresSchemaEditor  # type: ignore[assignment]
     introspection_class = PostgresIntrospection
     ops_class = PostgresOperations
 

psqlextra/backend/base_impl.py

@@ -3,6 +3,14 @@
 from django.conf import settings
 from django.core.exceptions import ImproperlyConfigured
 from django.db import DEFAULT_DB_ALIAS, connections
+from django.db.backends.postgresql.base import DatabaseWrapper
+from django.db.backends.postgresql.introspection import (  # type: ignore[import]
+    DatabaseIntrospection,
+)
+from django.db.backends.postgresql.operations import DatabaseOperations
+from django.db.backends.postgresql.schema import (  # type: ignore[import]
+    DatabaseSchemaEditor,
+)
 
 from django.db.backends.postgresql.base import (  # isort:skip
     DatabaseWrapper as Psycopg2DatabaseWrapper,
@@ -68,13 +76,13 @@ def base_backend_instance():
     return base_instance
 
 
-def backend():
+def backend() -> DatabaseWrapper:
     """Gets the base class for the database back-end."""
 
     return base_backend_instance().__class__
 
 
-def schema_editor():
+def schema_editor() -> DatabaseSchemaEditor:
     """Gets the base class for the schema editor.
 
     We have to use the configured base back-end's schema editor for
@@ -84,7 +92,7 @@ def schema_editor():
     return base_backend_instance().SchemaEditorClass
 
 
-def introspection():
+def introspection() -> DatabaseIntrospection:
     """Gets the base class for the introspection class.
 
     We have to use the configured base back-end's introspection class
@@ -94,7 +102,7 @@ def introspection():
     return base_backend_instance().introspection.__class__
 
 
-def operations():
+def operations() -> DatabaseOperations:
     """Gets the base class for the operations class.
 
     We have to use the configured base back-end's operations class for

psqlextra/backend/introspection.py

@@ -1,5 +1,9 @@
 from dataclasses import dataclass
-from typing import Dict, List, Optional, Tuple
+from typing import TYPE_CHECKING, Dict, List, Optional, Tuple
+
+from django.db.backends.postgresql.introspection import (  # type: ignore[import]
+    DatabaseIntrospection,
+)
 
 from psqlextra.types import PostgresPartitioningMethod
 
@@ -45,7 +49,16 @@ def partition_by_name(
         )
 
 
-class PostgresIntrospection(base_impl.introspection()):
+if TYPE_CHECKING:
+
+    class Introspection(DatabaseIntrospection):
+        pass
+
+else:
+    Introspection = base_impl.introspection()
+
+
+class PostgresIntrospection(Introspection):
     """Adds introspection features specific to PostgreSQL."""
 
     # TODO: This class is a mess, both here and in the
@@ -66,7 +79,7 @@ class PostgresIntrospection(base_impl.introspection()):
 
     def get_partitioned_tables(
         self, cursor
-    ) -> PostgresIntrospectedPartitonedTable:
+    ) -> List[PostgresIntrospectedPartitonedTable]:
         """Gets a list of partitioned tables."""
 
         cursor.execute(

psqlextra/backend/migrations/patched_autodetector.py

@@ -12,7 +12,7 @@
     RenameField,
 )
 from django.db.migrations.autodetector import MigrationAutodetector
-from django.db.migrations.operations.base import Operation
+from django.db.migrations.operations.fields import FieldOperation
 
 from psqlextra.models import (
     PostgresMaterializedViewModel,
@@ -83,7 +83,7 @@ def rename_field(self, operation: RenameField):
 
         return self._transform_view_field_operations(operation)
 
-    def _transform_view_field_operations(self, operation: Operation):
+    def _transform_view_field_operations(self, operation: FieldOperation):
         """Transforms operations on fields on a (materialized) view into state
         only operations.
 
@@ -199,9 +199,15 @@ def add_create_partitioned_model(self, operation: CreateModel):
                 )
             )
 
+        partitioned_kwargs = {
+            **kwargs,
+            "partitioning_options": partitioning_options,
+        }
+
         self.add(
             operations.PostgresCreatePartitionedModel(
-                *args, **kwargs, partitioning_options=partitioning_options
+                *args,
+                **partitioned_kwargs,
             )
         )
 
@@ -231,11 +237,9 @@ def add_create_view_model(self, operation: CreateModel):
 
         _, args, kwargs = operation.deconstruct()
 
-        self.add(
-            operations.PostgresCreateViewModel(
-                *args, **kwargs, view_options=view_options
-            )
-        )
+        view_kwargs = {**kwargs, "view_options": view_options}
+
+        self.add(operations.PostgresCreateViewModel(*args, **view_kwargs))
 
     def add_delete_view_model(self, operation: DeleteModel):
         """Adds a :see:PostgresDeleteViewModel operation to the list of
@@ -261,9 +265,12 @@ def add_create_materialized_view_model(self, operation: CreateModel):
 
         _, args, kwargs = operation.deconstruct()
 
+        view_kwargs = {**kwargs, "view_options": view_options}
+
         self.add(
             operations.PostgresCreateMaterializedViewModel(
-                *args, **kwargs, view_options=view_options
+                *args,
+                **view_kwargs,
             )
         )
 

psqlextra/backend/migrations/state/model.py

@@ -1,5 +1,5 @@
 from collections.abc import Mapping
-from typing import Type
+from typing import Tuple, Type, cast
 
 from django.db.migrations.state import ModelState
 from django.db.models import Model
@@ -17,8 +17,8 @@ class PostgresModelState(ModelState):
     """
 
     @classmethod
-    def from_model(
-        cls, model: PostgresModel, *args, **kwargs
+    def from_model(  # type: ignore[override]
+        cls, model: Type[PostgresModel], *args, **kwargs
     ) -> "PostgresModelState":
         """Creates a new :see:PostgresModelState object from the specified
         model.
@@ -29,28 +29,32 @@ def from_model(
         We also need to patch up the base class for the model.
         """
 
-        model_state = super().from_model(model, *args, **kwargs)
-        model_state = cls._pre_new(model, model_state)
+        model_state = super().from_model(
+            cast(Type[Model], model), *args, **kwargs
+        )
+        model_state = cls._pre_new(
+            model, cast("PostgresModelState", model_state)
+        )
 
         # django does not add abstract bases as a base in migrations
         # because it assumes the base does not add anything important
         # in a migration.. but it does, so we replace the Model
         # base with the actual base
-        bases = tuple()
+        bases: Tuple[Type[Model], ...] = tuple()
         for base in model_state.bases:
             if issubclass(base, Model):
                 bases += (cls._get_base_model_class(),)
             else:
                 bases += (base,)
 
-        model_state.bases = bases
+        model_state.bases = cast(Tuple[Type[Model]], bases)
         return model_state
 
     def clone(self) -> "PostgresModelState":
         """Gets an exact copy of this :see:PostgresModelState."""
 
         model_state = super().clone()
-        return self._pre_clone(model_state)
+        return self._pre_clone(cast(PostgresModelState, model_state))
 
     def render(self, apps):
         """Renders this state into an actual model."""
@@ -95,7 +99,9 @@ def render(self, apps):
 
     @classmethod
     def _pre_new(
-        cls, model: PostgresModel, model_state: "PostgresModelState"
+        cls,
+        model: Type[PostgresModel],
+        model_state: "PostgresModelState",
     ) -> "PostgresModelState":
         """Called when a new model state is created from the specified
         model."""

psqlextra/backend/migrations/state/partitioning.py

@@ -94,7 +94,7 @@ def delete_partition(self, name: str):
         del self.partitions[name]
 
     @classmethod
-    def _pre_new(
+    def _pre_new(  # type: ignore[override]
         cls,
         model: PostgresPartitionedModel,
         model_state: "PostgresPartitionedModelState",
@@ -108,7 +108,7 @@ def _pre_new(
         )
         return model_state
 
-    def _pre_clone(
+    def _pre_clone(  # type: ignore[override]
         self, model_state: "PostgresPartitionedModelState"
     ) -> "PostgresPartitionedModelState":
         """Called when this model state is cloned."""

psqlextra/backend/migrations/state/view.py

@@ -22,16 +22,18 @@ def __init__(self, *args, view_options={}, **kwargs):
         self.view_options = dict(view_options)
 
     @classmethod
-    def _pre_new(
-        cls, model: PostgresViewModel, model_state: "PostgresViewModelState"
+    def _pre_new(  # type: ignore[override]
+        cls,
+        model: Type[PostgresViewModel],
+        model_state: "PostgresViewModelState",
     ) -> "PostgresViewModelState":
         """Called when a new model state is created from the specified
         model."""
 
         model_state.view_options = dict(model._view_meta.original_attrs)
         return model_state
 
-    def _pre_clone(
+    def _pre_clone(  # type: ignore[override]
         self, model_state: "PostgresViewModelState"
     ) -> "PostgresViewModelState":
         """Called when this model state is cloned."""

psqlextra/backend/operations.py

@@ -9,7 +9,7 @@
 from . import base_impl
 
 
-class PostgresOperations(base_impl.operations()):
+class PostgresOperations(base_impl.operations()):  # type: ignore[misc]
     """Simple operations specific to PostgreSQL."""
 
     compiler_module = "psqlextra.compiler"