Guest User

Untitled

a guest
Mar 22nd, 2018
89
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
text 4.18 KB | None | 0 0
  1. import functools
  2. from django.db import models
  3. from django.db.models.constants import LOOKUP_SEP
  4. from django.db.models.query import ModelIterable
  5. from polymorphic.models import PolymorphicModel
  6.  
  7.  
  8. def rgetattr(obj, attr, separator='.'):
  9. """Recursive getattr to retrieve attributes of nested objects."""
  10. return functools.reduce(getattr, [obj] + attr.split(separator))
  11.  
  12.  
  13. def rsetattr(obj, attr, value, separator='.'):
  14. """Recursively getattr to fetch final object layer before using setattr."""
  15. attrs = attr.split(separator)
  16. setattr(functools.reduce(getattr, attrs[:-1], obj), attrs[-1], value)
  17.  
  18. def is_polymorphic_subclass(super_cls, sub_cls):
  19. try:
  20. return (issubclass(sub_cls, models.Model) and
  21. sub_cls != models.Model and
  22. sub_cls != super_cls and
  23. sub_cls != PolymorphicModel and
  24. super_cls in sub_cls._meta.parents)
  25. except AttributeError:
  26. return False
  27.  
  28.  
  29. def polymorphic_iterator(*fields):
  30. class PolymorphicModelIterable:
  31. def __init__(self, *args, **kwargs):
  32. self.iterable = ModelIterable(*args, **kwargs)
  33.  
  34. def __iter__(self):
  35. for obj in self.iterable:
  36. for field in fields:
  37. # Must get recursively in case our related polymorphic model is nested.
  38. instance = rgetattr(obj, field, separator=LOOKUP_SEP)
  39. real_instance_name = instance.polymorphic_ctype.model
  40. # We must copy the field cache for the base_model instance to the real instance
  41. # else additional data from select_related will be lost.
  42. real_instance = instance._state.fields_cache.pop(real_instance_name)
  43. real_instance._state.fields_cache = instance._state.fields_cache
  44. # Same recursion goes here for setting the related object.
  45. rsetattr(obj, field, real_instance, separator=LOOKUP_SEP)
  46. yield obj
  47.  
  48. return PolymorphicModelIterable
  49.  
  50.  
  51. class PolymorphicRelatedQuerySetMixin:
  52. """A class with a relationship to a polymorphic model should use this queryset"""
  53. def _get_nested_base_model(self, field):
  54. field_parts = field.split(LOOKUP_SEP)
  55. model = self.model
  56. for part in field_parts:
  57. field = getattr(model, part)
  58. # Should find a better solution to determine the related model than below.
  59. try:
  60. # In case of forward related descriptors.
  61. model = field.field.related_model
  62. except AttributeError:
  63. # In case of reverse related descriptors.
  64. model = field.related.related_model
  65.  
  66. return model
  67.  
  68. def select_polymorphic_related(self, *fields):
  69. """
  70. Specify fields that should be cast to the real polymorphic class.
  71. """
  72. subclass_names = []
  73. if fields:
  74. for field in fields:
  75. field_class = self._get_nested_base_model(field)
  76. # This is somewhat a replication of PolymorphicModel._get_inheritance_relation_fields_and_models(),
  77. # but it's necessary to do this unless we want to instantiate a base_model instance for every
  78. # query made. Would be a consideration to perhaps cache the results.
  79. for sub_cls in field_class.__subclasses__():
  80. if is_polymorphic_subclass(field_class, sub_cls):
  81. if sub_cls._meta.parents[field_class].remote_field.related_name:
  82. subclass_names.append(sub_cls._meta.parents[field_class].remote_field.related_name)
  83. else:
  84. subclass_names.append('{}__{}'.format(field, sub_cls.__name__.lower()))
  85. # We also need to add the polymorphic_ctype field name
  86. polymorphic_ctype_field_name = field_class.polymorphic_internal_model_fields[0]
  87. subclass_names.append('{}__{}'.format(field, polymorphic_ctype_field_name))
  88. self._iterable_class = polymorphic_iterator(*fields)
  89. return self.select_related(*subclass_names)
  90.  
  91.  
  92. class PolymorphicRelatedQuerySet(PolymorphicRelatedQuerySetMixin, models.QuerySet):
  93. pass
Add Comment
Please, Sign In to add comment