querysets.py 1.3 KB

123456789101112131415161718192021222324252627282930313233343536
  1. from django.db.models import QuerySet
  2. class PrefixQuerySet(QuerySet):
  3. def annotate_depth(self, limit=None):
  4. """
  5. Iterate through a QuerySet of Prefixes and annotate the hierarchical level of each. While it would be preferable
  6. to do this using .extra() on the QuerySet to count the unique parents of each prefix, that approach introduces
  7. performance issues at scale.
  8. Because we're adding a non-field attribute to the model, annotation must be made *after* any QuerySet
  9. modifications.
  10. """
  11. queryset = self
  12. stack = []
  13. for p in queryset:
  14. try:
  15. prev_p = stack[-1]
  16. except IndexError:
  17. prev_p = None
  18. if prev_p is not None:
  19. while (p.prefix not in prev_p.prefix) or p.prefix == prev_p.prefix:
  20. stack.pop()
  21. try:
  22. prev_p = stack[-1]
  23. except IndexError:
  24. prev_p = None
  25. break
  26. if prev_p is not None:
  27. prev_p.has_children = True
  28. stack.append(p)
  29. p.depth = len(stack) - 1
  30. if limit is None:
  31. return queryset
  32. return list(filter(lambda p: p.depth <= limit, queryset))