1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
|
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
|
