forked from tolomea/django-auto-prefetch
-
Notifications
You must be signed in to change notification settings - Fork 0
/
__init__.py
99 lines (71 loc) · 2.89 KB
/
__init__.py
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
from weakref import WeakValueDictionary
from django.db import models
from django.db.models.fields import related_descriptors
class DescriptorMixin:
def _field_name(self):
return self.field.name
def _is_cached(self, instance):
return self.is_cached(instance)
def _should_prefetch(self, instance):
return (
instance is not None # getattr on the class passes None to the descriptor
and not self._is_cached(instance) # already loaded
and len(getattr(instance, "_peers", [])) >= 2 # no peers no prefetch
)
def __get__(self, instance, instance_type=None):
if self._should_prefetch(instance):
prefetch = models.query.Prefetch(self._field_name())
peers = [p for p in instance._peers.values() if not self._is_cached(p)]
models.query.prefetch_related_objects(peers, prefetch)
return super().__get__(instance, instance_type)
class ForwardDescriptorMixin(DescriptorMixin):
def _should_prefetch(self, instance):
return super()._should_prefetch(
instance
) and None not in self.field.get_local_related_value(
instance
) # field is null
class ForwardManyToOneDescriptor(
ForwardDescriptorMixin, related_descriptors.ForwardManyToOneDescriptor
):
pass
class ForwardOneToOneDescriptor(
ForwardDescriptorMixin, related_descriptors.ForwardOneToOneDescriptor
):
pass
class ReverseOneToOneDescriptor(
DescriptorMixin, related_descriptors.ReverseOneToOneDescriptor
):
def _is_cached(self, instance):
return self.related.is_cached(instance)
def _field_name(self):
return self.related.get_accessor_name()
class ForeignKey(models.ForeignKey):
forward_related_accessor_class = ForwardManyToOneDescriptor
class OneToOneField(models.OneToOneField):
forward_related_accessor_class = ForwardOneToOneDescriptor
related_accessor_class = ReverseOneToOneDescriptor
class QuerySet(models.QuerySet):
def _fetch_all(self):
set_peers = self._result_cache is None
super()._fetch_all()
# ModelIterable tests for query sets returning model instances vs values or value lists etc
if set_peers and issubclass(self._iterable_class, models.query.ModelIterable):
peers = WeakValueDictionary((id(o), o) for o in self._result_cache)
for peer in peers.values():
peer._peers = peers
Manager = models.Manager.from_queryset(QuerySet)
class Model(models.Model):
class Meta:
abstract = True
base_manager_name = "prefetch_manager"
prefetch_manager = Manager()
objects = Manager()
def __getstate__(self):
# drop the peers info when pickling etc
res = super().__getstate__()
if "_peers" not in res: # pragma: no cover
return res
res = dict(res)
del res["_peers"]
return res