Skip to content

Commit

Permalink
Type-check entire code base with mypy
Browse files Browse the repository at this point in the history
Lots of `# type: ignore` in the internals, but the types should
be decent enough for consumption.
  • Loading branch information
Photonios committed Apr 20, 2023
1 parent 784e7a1 commit 86bc48d
Show file tree
Hide file tree
Showing 33 changed files with 315 additions and 126 deletions.
2 changes: 1 addition & 1 deletion .circleci/config.yml
Original file line number Diff line number Diff line change
Expand Up @@ -119,7 +119,7 @@ jobs:
steps:
- checkout
- install-dependencies:
extra: analysis
extra: analysis, test
- run:
name: Verify
command: python setup.py verify
Expand Down
20 changes: 18 additions & 2 deletions psqlextra/backend/base.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,7 @@
import logging

from typing import TYPE_CHECKING

from django.conf import settings
from django.db import ProgrammingError

Expand All @@ -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

Expand Down
16 changes: 12 additions & 4 deletions psqlextra/backend/base_impl.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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
Expand All @@ -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
Expand All @@ -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
Expand Down
19 changes: 16 additions & 3 deletions psqlextra/backend/introspection.py
Original file line number Diff line number Diff line change
@@ -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

Expand Down Expand Up @@ -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
Expand All @@ -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(
Expand Down
25 changes: 16 additions & 9 deletions psqlextra/backend/migrations/patched_autodetector.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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,
)
)

Expand Down Expand Up @@ -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
Expand All @@ -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,
)
)

Expand Down
24 changes: 15 additions & 9 deletions psqlextra/backend/migrations/state/model.py
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -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.
Expand All @@ -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."""
Expand Down Expand Up @@ -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."""
Expand Down
4 changes: 2 additions & 2 deletions psqlextra/backend/migrations/state/partitioning.py
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand All @@ -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."""
Expand Down
8 changes: 5 additions & 3 deletions psqlextra/backend/migrations/state/view.py
Original file line number Diff line number Diff line change
Expand Up @@ -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."""
Expand Down
2 changes: 1 addition & 1 deletion psqlextra/backend/operations.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down
Loading

0 comments on commit 86bc48d

Please sign in to comment.