src/catalogue/utils.py

   1 # This file is part of FNP-Redakcja, licensed under GNU Affero GPLv3 or later.
   2 # Copyright © Fundacja Nowoczesna Polska. See NOTICE for more information.
   3 #
   4 from collections import defaultdict
   5 from django.db.models import QuerySet
   6 from django.db.models.manager import BaseManager
   7
   8
   9 class UnrelatedQuerySet(QuerySet):
  10     def __init__(self, *args, **kwargs):
  11         super().__init__(*args, **kwargs)
  12         self._prefetch_unrelated_lookups = {}
  13         self._prefetch_unrelated_done = False
  14
  15     def _clone(self):
  16         c = super()._clone()
  17         c._prefetch_unrelated_lookups = self._prefetch_unrelated_lookups.copy()
  18         return c
  19
  20     def prefetch_unrelated(self, attribute, field, other_model, other_field):
  21         clone = self._clone()
  22         clone._prefetch_unrelated_lookups[field] = (attribute, other_model, other_field)
  23         return clone
  24
  25     def _fetch_all(self):
  26         prefetch_done = self._prefetch_done
  27         super()._fetch_all()
  28         if self._prefetch_unrelated_lookups and not prefetch_done:
  29             self._prefetch_unrelated_objects()
  30
  31     def _prefetch_unrelated_objects(self):
  32         for (
  33             field,
  34             (attribute, other_model, other_field),
  35         ) in self._prefetch_unrelated_lookups.items():
  36             values = set([getattr(obj, field) for obj in self._result_cache])
  37             other_objects = other_model._default_manager.filter(
  38                 **{f"{other_field}__in": values}
  39             )
  40             results = defaultdict(list)
  41             for other_obj in other_objects:
  42                 results[getattr(other_obj, other_field)].append(other_obj)
  43             for obj in self._result_cache:
  44                 setattr(obj, attribute, results.get(getattr(obj, field)))
  45         self._prefetch_unrelated_done = True
  46
  47
  48 class UnrelatedManager(BaseManager.from_queryset(UnrelatedQuerySet)):
  49     pass