/polymorphic/query.py

  1# -*- coding: utf-8 -*-
  2""" QuerySet for PolymorphicModel
  3    Please see README.rst or DOCS.rst or http://bserve.webhop.org/wiki/django_polymorphic
  4"""
  5
  6from compatibility_tools import defaultdict
  7
  8from django.db.models.query import QuerySet
  9from django.contrib.contenttypes.models import ContentType
 10
 11from query_translate import translate_polymorphic_filter_definitions_in_kwargs, translate_polymorphic_filter_definitions_in_args
 12from query_translate import translate_polymorphic_field_path
 13
 14# chunk-size: maximum number of objects requested per db-request
 15# by the polymorphic queryset.iterator() implementation; we use the same chunk size as Django
 16from django.db.models.query import CHUNK_SIZE               # this is 100 for Django 1.1/1.2
 17Polymorphic_QuerySet_objects_per_request = CHUNK_SIZE
 18
 19
 20
 21###################################################################################
 22### PolymorphicQuerySet
 23
 24class PolymorphicQuerySet(QuerySet):
 25    """
 26    QuerySet for PolymorphicModel
 27
 28    Contains the core functionality for PolymorphicModel
 29
 30    Usually not explicitly needed, except if a custom queryset class
 31    is to be used.
 32    """
 33
 34    def __init__(self, *args, **kwargs):
 35        "init our queryset object member variables"
 36        self.polymorphic_disabled = False
 37        super(PolymorphicQuerySet, self).__init__(*args, **kwargs)
 38
 39    def _clone(self, *args, **kwargs):
 40        "Django's _clone only copies its own variables, so we need to copy ours here"
 41        new = super(PolymorphicQuerySet, self)._clone(*args, **kwargs)
 42        new.polymorphic_disabled = self.polymorphic_disabled
 43        return new
 44
 45    def non_polymorphic(self, *args, **kwargs):
 46        """switch off polymorphic behaviour for this query.
 47        When the queryset is evaluated, only objects of the type of the
 48        base class used for this query are returned."""
 49        self.polymorphic_disabled = True
 50        return self
 51
 52    def instance_of(self, *args):
 53        """Filter the queryset to only include the classes in args (and their subclasses).
 54        Implementation in _translate_polymorphic_filter_defnition."""
 55        return self.filter(instance_of=args)
 56
 57    def not_instance_of(self, *args):
 58        """Filter the queryset to exclude the classes in args (and their subclasses).
 59        Implementation in _translate_polymorphic_filter_defnition."""
 60        return self.filter(not_instance_of=args)
 61
 62    def _filter_or_exclude(self, negate, *args, **kwargs):
 63        "We override this internal Django functon as it is used for all filter member functions."
 64        translate_polymorphic_filter_definitions_in_args(self.model, args) # the Q objects
 65        additional_args = translate_polymorphic_filter_definitions_in_kwargs(self.model, kwargs) # filter_field='data'
 66        return super(PolymorphicQuerySet, self)._filter_or_exclude(negate, *(list(args) + additional_args), **kwargs)
 67
 68    def order_by(self, *args, **kwargs):
 69        """translate the field paths in the args, then call vanilla order_by."""
 70        new_args = [ translate_polymorphic_field_path(self.model, a) for a in args ]
 71        return super(PolymorphicQuerySet, self).order_by(*new_args, **kwargs)
 72
 73    def _process_aggregate_args(self, args, kwargs):
 74        """for aggregate and annotate kwargs: allow ModelX___field syntax for kwargs, forbid it for args.
 75        Modifies kwargs if needed (these are Aggregate objects, we translate the lookup member variable)"""
 76        for a in args:
 77            assert not '___' in a.lookup, 'PolymorphicModel: annotate()/aggregate(): ___ model lookup supported for keyword arguments only'
 78        for a in kwargs.values():
 79            a.lookup = translate_polymorphic_field_path(self.model, a.lookup)
 80
 81    def annotate(self, *args, **kwargs):
 82        """translate the polymorphic field paths in the kwargs, then call vanilla annotate.
 83        _get_real_instances will do the rest of the job after executing the query."""
 84        self._process_aggregate_args(args, kwargs)
 85        return super(PolymorphicQuerySet, self).annotate(*args, **kwargs)
 86
 87    def aggregate(self, *args, **kwargs):
 88        """translate the polymorphic field paths in the kwargs, then call vanilla aggregate.
 89        We need no polymorphic object retrieval for aggregate => switch it off."""
 90        self._process_aggregate_args(args, kwargs)
 91        self.polymorphic_disabled = True
 92        return super(PolymorphicQuerySet, self).aggregate(*args, **kwargs)
 93
 94    # Since django_polymorphic 'V1.0 beta2', extra() always returns polymorphic results.^
 95    # The resulting objects are required to have a unique primary key within the result set
 96    # (otherwise an error is thrown).
 97    # The "polymorphic" keyword argument is not supported anymore.
 98    #def extra(self, *args, **kwargs):
 99
100
101    def _get_real_instances(self, base_result_objects):
102        """
103        Polymorphic object loader
104
105        Does the same as:
106
107            return [ o.get_real_instance() for o in base_result_objects ]
108
109        but more efficiently.
110
111        The list base_result_objects contains the objects from the executed
112        base class query. The class of all of them is self.model (our base model).
113
114        Some, many or all of these objects were not created and stored as
115        class self.model, but as a class derived from self.model. We want to re-fetch
116        these objects from the db as their original class so we can return them
117        just as they were created/saved.
118
119        We identify these objects by looking at o.polymorphic_ctype, which specifies
120        the real class of these objects (the class at the time they were saved).
121
122        First, we sort the result objects in base_result_objects for their
123        subclass (from o.polymorphic_ctype), and then we execute one db query per
124        subclass of objects. Here, we handle any annotations from annotate().
125
126        Finally we re-sort the resulting objects into the correct order and
127        return them as a list.
128        """
129        ordered_id_list = []    # list of ids of result-objects in correct order
130        results = {}            # polymorphic dict of result-objects, keyed with their id (no order)
131
132        # dict contains one entry per unique model type occurring in result,
133        # in the format idlist_per_model[modelclass]=[list-of-object-ids]
134        idlist_per_model = defaultdict(list)
135
136        # - sort base_result_object ids into idlist_per_model lists, depending on their real class;
137        # - also record the correct result order in "ordered_id_list"
138        # - store objects that already have the correct class into "results"
139        base_result_objects_by_id = {}
140        self_model_content_type_id = ContentType.objects.get_for_model(self.model).pk
141        for base_object in base_result_objects:
142            ordered_id_list.append(base_object.pk)
143
144            # check if id of the result object occeres more than once - this can happen e.g. with base_objects.extra(tables=...)
145            assert not base_object.pk in base_result_objects_by_id, (
146                "django_polymorphic: result objects do not have unique primary keys - model "+unicode(self.model) )
147
148            base_result_objects_by_id[base_object.pk] = base_object
149
150            # this object is not a derived object and already the real instance => store it right away
151            if (base_object.polymorphic_ctype_id == self_model_content_type_id):
152                results[base_object.pk] = base_object
153
154            # this object is derived and its real instance needs to be retrieved
155            # => store it's id into the bin for this model type
156            else:
157                idlist_per_model[base_object.get_real_instance_class()].append(base_object.pk)
158
159        # django's automatic ".pk" field does not always work correctly for
160        # custom fields in derived objects (unclear yet who to put the blame on).
161        # We get different type(o.pk) in this case.
162        # We work around this by using the real name of the field directly
163        # for accessing the primary key of the the derived objects.
164        # We might assume that self.model._meta.pk.name gives us the name of the primary key field,
165        # but it doesn't. Therefore we use polymorphic_primary_key_name, which we set up in base.py.
166        pk_name = self.model.polymorphic_primary_key_name
167
168        # For each model in "idlist_per_model" request its objects (the real model)
169        # from the db and store them in results[].
170        # Then we copy the annotate fields from the base objects to the real objects.
171        # Then we copy the extra() select fields from the base objects to the real objects.
172        # TODO: defer(), only(): support for these would be around here
173        for modelclass, idlist in idlist_per_model.items():
174            qs = modelclass.base_objects.filter(pk__in=idlist) # use pk__in instead ####
175            qs.dup_select_related(self)    # copy select related configuration to new qs
176
177            for o in qs:
178                o_pk=getattr(o,pk_name)
179                
180                if self.query.aggregates:
181                    for anno_field_name in self.query.aggregates.keys():
182                        attr = getattr(base_result_objects_by_id[o_pk], anno_field_name)
183                        setattr(o, anno_field_name, attr)
184
185                if self.query.extra_select:
186                    for select_field_name in self.query.extra_select.keys():
187                        attr = getattr(base_result_objects_by_id[o_pk], select_field_name)
188                        setattr(o, select_field_name, attr)
189
190                results[o_pk] = o
191
192        # re-create correct order and return result list
193        resultlist = [ results[ordered_id] for ordered_id in ordered_id_list if ordered_id in results ]
194
195        # set polymorphic_annotate_names in all objects (currently just used for debugging/printing)
196        if self.query.aggregates:
197            annotate_names=self.query.aggregates.keys() # get annotate field list
198            for o in resultlist:
199                o.polymorphic_annotate_names=annotate_names
200
201        # set polymorphic_extra_select_names in all objects (currently just used for debugging/printing)
202        if self.query.extra_select:
203            extra_select_names=self.query.extra_select.keys() # get extra select field list
204            for o in resultlist:
205                o.polymorphic_extra_select_names=extra_select_names
206            
207        return resultlist
208
209    def iterator(self):
210        """
211        This function is used by Django for all object retrieval.
212        By overriding it, we modify the objects that this queryset returns
213        when it is evaluated (or its get method or other object-returning methods are called).
214
215        Here we do the same as:
216
217            base_result_objects=list(super(PolymorphicQuerySet, self).iterator())
218            real_results=self._get_real_instances(base_result_objects)
219            for o in real_results: yield o
220
221        but it requests the objects in chunks from the database,
222        with Polymorphic_QuerySet_objects_per_request per chunk
223        """
224        base_iter = super(PolymorphicQuerySet, self).iterator()
225
226        # disabled => work just like a normal queryset
227        if self.polymorphic_disabled:
228            for o in base_iter: yield o
229            raise StopIteration
230
231        while True:
232            base_result_objects = []
233            reached_end = False
234
235            for i in range(Polymorphic_QuerySet_objects_per_request):
236                try:
237                    o=base_iter.next()
238                    base_result_objects.append(o)
239                except StopIteration:
240                    reached_end = True
241                    break
242
243            real_results = self._get_real_instances(base_result_objects)
244
245            for o in real_results:
246                yield o
247
248            if reached_end: raise StopIteration
249
250    def __repr__(self, *args, **kwargs):
251        if self.model.polymorphic_query_multiline_output:
252            result = [ repr(o) for o in self.all() ]
253            return  '[ ' + ',\n  '.join(result) + ' ]'
254        else:
255            return super(PolymorphicQuerySet,self).__repr__(*args, **kwargs)
256
257    class _p_list_class(list):
258        def __repr__(self, *args, **kwargs):
259            result = [ repr(o) for o in self ]
260            return  '[ ' + ',\n  '.join(result) + ' ]'
261
262    def get_real_instances(self, base_result_objects=None):
263        "same as _get_real_instances, but make sure that __repr__ for ShowField... creates correct output"
264        if not base_result_objects: base_result_objects=self
265        olist = self._get_real_instances(base_result_objects)
266        if not self.model.polymorphic_query_multiline_output:
267            return olist
268        clist=PolymorphicQuerySet._p_list_class(olist)
269        return clist
270