Coverage for utilities/fields.py: 32%
183 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-10-10 18:35 +0000
1from collections import defaultdict
3from django.contrib.contenttypes.fields import GenericForeignKey, GenericForeignKeyDescriptor
4from django.contrib.contenttypes.models import ContentType
5from django.core.exceptions import ObjectDoesNotExist
6from django.db import models
7from django.db.models.fields.mixins import FieldCacheMixin
8from django.utils.functional import cached_property
9from django.utils.translation import gettext_lazy as _
11from .forms.widgets import ColorSelect
12from .validators import ColorValidator
14__all__ = (
15 'ColorField',
16 'CounterCacheField',
17 'GenericArrayForeignKey',
18 'NaturalOrderingField',
19 'RestrictedGenericForeignKey',
20)
23class ColorField(models.CharField):
24 default_validators = [ColorValidator]
25 description = "A hexadecimal RGB color code"
27 def __init__(self, *args, **kwargs):
28 kwargs['max_length'] = 6
29 super().__init__(*args, **kwargs)
31 def formfield(self, **kwargs):
32 kwargs['widget'] = ColorSelect
33 return super().formfield(**kwargs)
36class NaturalOrderingField(models.CharField):
37 """
38 A field which stores a naturalized representation of its target field, to be used for ordering its parent model.
40 :param target_field: Name of the field of the parent model to be naturalized
41 :param naturalize_function: The function used to generate a naturalized value (optional)
42 """
43 description = "Stores a representation of its target field suitable for natural ordering"
45 def __init__(self, target_field, naturalize_function, *args, **kwargs):
46 self.target_field = target_field
47 self.naturalize_function = naturalize_function
48 super().__init__(*args, **kwargs)
50 def pre_save(self, model_instance, add):
51 """
52 Generate a naturalized value from the target field
53 """
54 original_value = getattr(model_instance, self.target_field)
55 naturalized_value = self.naturalize_function(original_value, max_length=self.max_length)
56 setattr(model_instance, self.attname, naturalized_value)
58 return naturalized_value
60 def deconstruct(self):
61 kwargs = super().deconstruct()[3] # Pass kwargs from CharField
62 kwargs['naturalize_function'] = self.naturalize_function
63 return (
64 self.name,
65 'utilities.fields.NaturalOrderingField',
66 [self.target_field],
67 kwargs,
68 )
71class RestrictedGenericForeignKeyDescriptor(GenericForeignKeyDescriptor):
72 """
73 Django 6.1 moved get_prefetch_querysets() off GenericForeignKey and onto a separate
74 descriptor, which prefetch_related() now prefers over the field itself. Delegate back to
75 the field so that RestrictedGenericForeignKey's restrict()-aware implementation is used.
76 """
77 def get_prefetch_querysets(self, instances, querysets=None):
78 return self.field.get_prefetch_querysets(instances, querysets)
81class RestrictedGenericForeignKey(GenericForeignKey):
83 def contribute_to_class(self, cls, name, **kwargs):
84 super().contribute_to_class(cls, name, **kwargs)
85 # Replace the descriptor installed by GenericForeignKey with one which defers to
86 # get_prefetch_querysets() below.
87 setattr(cls, self.attname, RestrictedGenericForeignKeyDescriptor(self))
89 # Replicated largely from GenericForeignKeyDescriptor. Changes include:
90 # 1. Capture restrict_params from RestrictedPrefetch (hack)
91 # 2. If restrict_params is set, call restrict() on the queryset for
92 # the related model
93 def get_prefetch_querysets(self, instances, querysets=None):
94 restrict_params = {}
95 custom_queryset_dict = {}
97 # Compensate for the hack in RestrictedPrefetch
98 if type(querysets) is dict:
99 restrict_params = querysets
101 elif querysets is not None:
102 for queryset in querysets:
103 ct_id = self.get_content_type(
104 model=queryset.query.model, using=queryset.db
105 ).pk
106 if ct_id in custom_queryset_dict:
107 raise ValueError(
108 "Only one queryset is allowed for each content type."
109 )
110 custom_queryset_dict[ct_id] = queryset
112 # For efficiency, group the instances by content type and then do one
113 # query per model
114 fk_dict = defaultdict(set)
115 # We need one instance for each group in order to get the right db:
116 instance_dict = {}
117 ct_attname = self.model._meta.get_field(self.ct_field).get_attname()
118 for instance in instances:
119 # We avoid looking for values if either ct_id or fkey value is None
120 ct_id = getattr(instance, ct_attname)
121 if ct_id is not None:
122 # Check if the content type actually exists
123 if not self.get_content_type(id=ct_id, using=instance._state.db).model_class():
124 continue
126 fk_val = getattr(instance, self.fk_field)
127 if fk_val is not None:
128 fk_dict[ct_id].add(fk_val)
129 instance_dict[ct_id] = instance
131 ret_val = []
132 for ct_id, fkeys in fk_dict.items():
133 if ct_id in custom_queryset_dict:
134 # Return values from the custom queryset, if provided.
135 qs = custom_queryset_dict[ct_id].filter(pk__in=fkeys)
136 else:
137 instance = instance_dict[ct_id]
138 ct = self.get_content_type(id=ct_id, using=instance._state.db)
139 model = ct.model_class()
140 # A model may declare its manager under a name other than `objects`
141 manager = getattr(model, 'objects', model._default_manager)
142 qs = manager.filter(pk__in=fkeys)
143 if restrict_params:
144 qs = qs.restrict(**restrict_params)
145 # Carry the fetch mode of the objects being prefetched over to the objects prefetched
146 # onto them. Every instance in a batch shares one fetch mode, so the first is
147 # representative; it is safe to index because fk_dict is populated from `instances`,
148 # and so is empty (skipping this loop entirely) whenever `instances` is.
149 ret_val.extend(qs.fetch_mode(instances[0]._state.fetch_mode))
151 # For doing the join in Python, we have to match both the FK val and the
152 # content type, so we use a callable that returns a (fk, class) pair.
153 def gfk_key(obj):
154 ct_id = getattr(obj, ct_attname)
155 if ct_id is None:
156 return None
157 if model := self.get_content_type(
158 id=ct_id, using=obj._state.db
159 ).model_class():
160 return (
161 model._meta.pk.get_prep_value(getattr(obj, self.fk_field)),
162 model,
163 )
164 return None
166 return (
167 ret_val,
168 lambda obj: (obj.pk, obj.__class__),
169 gfk_key,
170 True,
171 self.name,
172 False,
173 )
176class CounterCacheField(models.BigIntegerField):
177 """
178 Counter field to keep track of related model counts.
179 """
180 def __init__(self, to_model, to_field, *args, **kwargs):
181 if not isinstance(to_model, str): 181 ↛ 182line 181 didn't jump to line 182 because the condition on line 181 was never true
182 raise TypeError(
183 _("%s(%r) is invalid. to_model parameter to CounterCacheField must be "
184 "a string in the format 'app.model'")
185 % (
186 self.__class__.__name__,
187 to_model,
188 )
189 )
191 if not isinstance(to_field, str): 191 ↛ 192line 191 didn't jump to line 192 because the condition on line 191 was never true
192 raise TypeError(
193 _("%s(%r) is invalid. to_field parameter to CounterCacheField must be "
194 "a string in the format 'field'")
195 % (
196 self.__class__.__name__,
197 to_field,
198 )
199 )
201 self.to_model_name = to_model
202 self.to_field_name = to_field
204 kwargs['default'] = kwargs.get('default', 0)
205 kwargs['editable'] = False
207 super().__init__(*args, **kwargs)
209 def deconstruct(self):
210 name, path, args, kwargs = super().deconstruct()
211 kwargs["to_model"] = self.to_model_name
212 kwargs["to_field"] = self.to_field_name
213 return name, path, args, kwargs
216class GenericArrayForeignKey(FieldCacheMixin, models.Field):
217 """
218 Provide a generic many-to-many relation through an 2d array field
219 """
221 many_to_many = False
222 many_to_one = False
223 one_to_many = True
224 one_to_one = False
226 def __init__(self, field, for_concrete_model=True):
227 super().__init__(editable=False)
228 self.field = field
229 self.for_concrete_model = for_concrete_model
230 self.is_relation = True
232 def contribute_to_class(self, cls, name, **kwargs):
233 super().contribute_to_class(cls, name, private_only=True, **kwargs)
234 # GenericArrayForeignKey is its own descriptor.
235 setattr(cls, self.attname, self)
237 @cached_property
238 def cache_name(self):
239 return self.name
241 def get_cache_name(self):
242 return self.cache_name
244 def _get_ids(self, instance):
245 return getattr(instance, self.field)
247 def get_content_type_by_id(self, id=None, using=None):
248 return ContentType.objects.db_manager(using).get_for_id(id)
250 def get_content_type_of_obj(self, obj=None):
251 return ContentType.objects.db_manager(obj._state.db).get_for_model(
252 obj, for_concrete_model=self.for_concrete_model
253 )
255 def get_content_type_for_model(self, using=None, model=None):
256 return ContentType.objects.db_manager(using).get_for_model(
257 model, for_concrete_model=self.for_concrete_model
258 )
260 def get_prefetch_querysets(self, instances, querysets=None):
261 custom_queryset_dict = {}
262 if querysets is not None:
263 for queryset in querysets:
264 ct_id = self.get_content_type_for_model(
265 model=queryset.query.model, using=queryset.db
266 ).pk
267 if ct_id in custom_queryset_dict:
268 raise ValueError(
269 "Only one queryset is allowed for each content type."
270 )
271 custom_queryset_dict[ct_id] = queryset
273 # For efficiency, group the instances by content type and then do one
274 # query per model
275 fk_dict = defaultdict(set) # type id, db -> model ids
276 for instance in instances:
277 for step in self._get_ids(instance):
278 for ct_id, fk_val in step:
279 fk_dict[(ct_id, instance._state.db)].add(fk_val)
281 rel_objects = []
282 for (ct_id, db), fkeys in fk_dict.items():
283 if ct_id in custom_queryset_dict:
284 rel_objects.extend(custom_queryset_dict[ct_id].filter(pk__in=fkeys))
285 else:
286 ct = self.get_content_type_by_id(id=ct_id, using=db)
287 rel_objects.extend(ct.get_all_objects_for_this_type(pk__in=fkeys))
289 # reorganize objects to fix usage
290 items = {
291 (self.get_content_type_of_obj(obj=rel_obj).pk, rel_obj.pk, rel_obj._state.db): rel_obj
292 for rel_obj in rel_objects
293 }
294 lists = []
295 lists_keys = {}
296 for instance in instances:
297 data = []
298 lists.append(data)
299 lists_keys[instance] = id(data)
300 for step in self._get_ids(instance):
301 nodes = []
302 for ct, fk in step:
303 if rel_obj := items.get((ct, fk, instance._state.db)):
304 nodes.append(rel_obj)
305 data.append(nodes)
307 return (
308 lists,
309 lambda obj: id(obj),
310 lambda obj: lists_keys[obj],
311 True,
312 self.cache_name,
313 False,
314 )
316 def __get__(self, instance, cls=None):
317 if instance is None:
318 return self
319 rel_objects = self.get_cached_value(instance, default=...)
320 expected_ids = self._get_ids(instance)
321 # we do not check if cache actual
322 if rel_objects is not ...:
323 return rel_objects
324 # load value
325 if expected_ids is None:
326 self.set_cached_value(instance, rel_objects)
327 return rel_objects
328 data = []
329 for step in self._get_ids(instance):
330 rel_objects = []
331 for ct_id, pk_val in step:
332 ct = self.get_content_type_by_id(id=ct_id, using=instance._state.db)
333 try:
334 rel_obj = ct.get_object_for_this_type(pk=pk_val)
335 rel_objects.append(rel_obj)
336 except ObjectDoesNotExist:
337 pass
338 data.append(rel_objects)
339 self.set_cached_value(instance, data)
340 return data