1import warnings
2from collections import defaultdict
3from collections.abc import Callable
4from functools import partial
5from typing import TYPE_CHECKING
6
7if TYPE_CHECKING:
8 from plain.postgres.base import Model
9
10
11class ModelsRegistryNotReady(Exception):
12 """The plain.postgres registry is not populated yet"""
13
14
15class ModelsRegistry:
16 def __init__(self) -> None:
17 # Memo for get_models(), keyed by package_label ("" = all). Cleared by
18 # clear_cache() whenever the registry changes.
19 self._get_models_cache: dict[str, list[type[Model]]] = {}
20
21 # Mapping of app labels => model names => model classes. Every time a
22 # model is imported, ModelBase.__new__ calls packages.register_model which
23 # creates an entry in all_models. All imported models are registered,
24 # regardless of whether they're defined in an installed application
25 # and whether the registry has been populated. Since it isn't possible
26 # to reimport a module safely (it could reexecute initialization code)
27 # all_models is never overridden or reset.
28 self.all_models: defaultdict[str, dict[str, type[Model]]] = defaultdict(dict)
29
30 # Maps ("package_label", "modelname") tuples to lists of functions to be
31 # called when the corresponding model is ready. Used by this class's
32 # `lazy_model_operation()` and `do_pending_operations()` methods.
33 self._pending_operations: defaultdict[
34 tuple[str, str], list[Callable[[type[Model]], None]]
35 ] = defaultdict(list)
36
37 self.ready: bool = False
38
39 def check_ready(self) -> None:
40 """Raise an exception if all models haven't been imported yet."""
41 if not self.ready:
42 raise ModelsRegistryNotReady("Models aren't loaded yet.")
43
44 # This method is performance-critical at least for Plain's test suite.
45 def get_models(self, *, package_label: str = "") -> list[type[Model]]:
46 """
47 Return a list of all installed models.
48
49 By default, the following models aren't included:
50
51 - auto-created models for many-to-many relations without
52 an explicit intermediate table,
53
54 Set the corresponding keyword argument to True to include such models.
55 """
56
57 self.check_ready()
58
59 if (cached := self._get_models_cache.get(package_label)) is not None:
60 return cached
61
62 models = []
63
64 # Get models for a single package
65 if package_label:
66 models.extend(self.all_models[package_label].values())
67 else:
68 # Get models for all packages
69 for package_models in self.all_models.values():
70 models.extend(package_models.values())
71
72 self._get_models_cache[package_label] = models
73 return models
74
75 def get_model(
76 self,
77 package_label: str,
78 model_name: str | None = None,
79 require_ready: bool = True,
80 ) -> type[Model]:
81 """
82 Return the model matching the given package_label and model_name.
83
84 As a shortcut, package_label may be in the form <package_label>.<model_name>.
85
86 model_name is case-insensitive.
87
88 Raise LookupError if no application exists with this label, or no
89 model exists with this name in the application. Raise ValueError if
90 called with a single argument that doesn't contain exactly one dot.
91 """
92
93 if require_ready:
94 self.check_ready()
95
96 if model_name is None:
97 package_label, model_name = package_label.split(".")
98
99 package_models = self.all_models[package_label]
100 return package_models[model_name.lower()]
101
102 def register_model(self, package_label: str, model: type[Model]) -> None:
103 # Since this method is called when models are imported, it cannot
104 # perform imports because of the risk of import loops. It mustn't
105 # call get_package_config().
106 model_name = model.model_options.model_name
107 app_models = self.all_models[package_label]
108 if model_name in app_models:
109 if (
110 model.__name__ == app_models[model_name].__name__
111 and model.__module__ == app_models[model_name].__module__
112 ):
113 warnings.warn(
114 f"Model '{package_label}.{model_name}' was already registered. Reloading models is not "
115 "advised as it can lead to inconsistencies, most notably with "
116 "related models.",
117 RuntimeWarning,
118 stacklevel=2,
119 )
120 else:
121 raise RuntimeError(
122 f"Conflicting '{model_name}' models in application '{package_label}': {app_models[model_name]} and {model}."
123 )
124 app_models[model_name] = model
125 self.do_pending_operations(model)
126 self.clear_cache()
127
128 def _get_registered_model(self, package_label: str, model_name: str) -> type[Model]:
129 """
130 Similar to get_model(), but doesn't require that an app exists with
131 the given package_label.
132
133 It's safe to call this method at import time, even while the registry
134 is being populated.
135 """
136 model = self.all_models[package_label].get(model_name.lower())
137 if model is None:
138 raise LookupError(f"Model '{package_label}.{model_name}' not registered.")
139 return model
140
141 def clear_cache(self) -> None:
142 """
143 Clear all internal caches, for methods that alter the app registry.
144
145 This is mostly used in tests.
146 """
147 # Call expire cache on each model. This will purge
148 # the relation tree and the fields cache.
149 self._get_models_cache.clear()
150 if self.ready:
151 # Circumvent self.get_models() to prevent that the cache is refilled.
152 # This particularly prevents that an empty value is cached while cloning.
153 for package_models in self.all_models.values():
154 for model in package_models.values():
155 model._model_meta._expire_cache()
156
157 def lazy_model_operation(
158 self, function: Callable[..., None], *model_keys: tuple[str, str]
159 ) -> None:
160 """
161 Take a function and a number of ("package_label", "modelname") tuples, and
162 when all the corresponding models have been imported and registered,
163 call the function with the model classes as its arguments.
164
165 The function passed to this method must accept exactly n models as
166 arguments, where n=len(model_keys).
167 """
168 # Base case: no arguments, just execute the function.
169 if not model_keys:
170 function()
171 # Recursive case: take the head of model_keys, wait for the
172 # corresponding model class to be imported and registered, then apply
173 # that argument to the supplied function. Pass the resulting partial
174 # to lazy_model_operation() along with the remaining model args and
175 # repeat until all models are loaded and all arguments are applied.
176 else:
177 next_model, *more_models = model_keys
178
179 # This will be executed after the class corresponding to next_model
180 # has been imported and registered.
181 def apply_next_model(model: type[Model]) -> None:
182 next_function = partial(function, model)
183 self.lazy_model_operation(next_function, *more_models)
184
185 # If the model has already been imported and registered, partially
186 # apply it to the function now. If not, add it to the list of
187 # pending operations for the model, where it will be executed with
188 # the model class as its sole argument once the model is ready.
189 try:
190 model_class = self._get_registered_model(*next_model)
191 except LookupError:
192 self._pending_operations[next_model].append(apply_next_model)
193 else:
194 apply_next_model(model_class)
195
196 def do_pending_operations(self, model: type[Model]) -> None:
197 """
198 Take a newly-prepared model and pass it to each function waiting for
199 it. This is called at the very end of Models.register_model().
200 """
201 key = model.model_options.package_label, model.model_options.model_name
202 for function in self._pending_operations.pop(key, []):
203 function(model)
204
205
206models_registry = ModelsRegistry()
207
208
209# Decorator to register a model (using the internal registry for the correct state).
210def register_model[M: "Model"](model_class: type[M]) -> type[M]:
211 model_class._model_meta.models_registry.register_model(
212 model_class.model_options.package_label,
213 model_class,
214 )
215 return model_class