/polymorphic/query.py

https://bitbucket.org/bconstantin/django_polymorphic/ · Python · 270 lines · 121 code · 48 blank · 101 comment · 32 complexity · 0c6eb8cbc4edcf15bc35a94bf7e3e7a4 MD5 · raw file

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