GlitchTip/psqlextra/backend/migrations/state/model.py

from collections.abc import Mapping
from typing import Tuple, Type, cast

from django.db.migrations.state import ModelState
from django.db.models import Model

from psqlextra.models import PostgresModel


class PostgresModelState(ModelState):
    """Base for custom model states.

    We need this base class to create some hooks into rendering models,
    creating new states and cloning state. Most of the logic resides
    here in the base class. Our derived classes implement the `_pre_*`
    methods.
    """

    @classmethod
    def from_model(  # type: ignore[override]
        cls, model: Type[PostgresModel], *args, **kwargs
    ) -> "PostgresModelState":
        """Creates a new :see:PostgresModelState object from the specified
        model.

        We override this so derived classes get the chance to attach
        additional information to the newly created model state.

        We also need to patch up the base class for the model.
        """

        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[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 = 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(cast(PostgresModelState, model_state))

    def render(self, apps):
        """Renders this state into an actual model."""

        # TODO: figure out a way to do this witout pretty much
        #       copying the base class's implementation

        try:
            bases = tuple(
                (apps.get_model(base) if isinstance(base, str) else base)
                for base in self.bases
            )
        except LookupError:
            # TODO: this should be a InvalidBaseError
            raise ValueError(
                "Cannot resolve one or more bases from %r" % (self.bases,)
            )

        if isinstance(self.fields, Mapping):
            # In Django 3.1 `self.fields` became a `dict`
            fields = {
                name: field.clone() for name, field in self.fields.items()
            }
        else:
            # In Django < 3.1 `self.fields` is a list of (name, field) tuples
            fields = {name: field.clone() for name, field in self.fields}

        meta = type(
            "Meta",
            (),
            {"app_label": self.app_label, "apps": apps, **self.options},
        )

        attributes = {
            **fields,
            "Meta": meta,
            "__module__": "__fake__",
            **dict(self.construct_managers()),
        }

        return type(*self._pre_render(self.name, bases, attributes))

    @classmethod
    def _pre_new(
        cls,
        model: Type[PostgresModel],
        model_state: "PostgresModelState",
    ) -> "PostgresModelState":
        """Called when a new model state is created from the specified
        model."""

        return model_state

    def _pre_clone(
        self, model_state: "PostgresModelState"
    ) -> "PostgresModelState":
        """Called when this model state is cloned."""

        return model_state

    def _pre_render(self, name: str, bases, attributes):
        """Called when this model state is rendered into a model."""

        return name, bases, attributes

    @classmethod
    def _get_base_model_class(self) -> Type[PostgresModel]:
        """Gets the class to use as a base class for rendered models."""

        return PostgresModel