diff --git a/emeis/core/models.py b/emeis/core/models.py index 00360be..7872d57 100644 --- a/emeis/core/models.py +++ b/emeis/core/models.py @@ -1,5 +1,7 @@ +import operator import unicodedata import uuid +from functools import reduce from django.conf import settings from django.contrib.auth.models import AbstractBaseUser, UserManager @@ -158,6 +160,53 @@ def is_authenticated(self): return True +class ScopeQuerySet(TreeQuerySet): + def descendants(self, include_self=False): + """Return a QS that contains all descendants of the given QS. + + This is a workaround for django-tree-queries, which currently does + not support this query (it can only do it on single nodes). + """ + # django-tree-queries sadly does not (yet?) support descendants query + # for QS - only for single nodes + descendants_q = reduce( + operator.or_, + [ + models.Q( + pk__in=entry.descendants(include_self=include_self).values_list( + "pk", flat=True + ) + ) + for entry in self + ], + initial=models.Q(), + ) + return self.model.objects.filter(descendants_q) + + def ancestors(self, include_self=False): + """Return a QS that contains all ancestors of the given QS. + + This is a workaround for django-tree-queries, which currently does + not support this query (it can only do it on single nodes). + """ + # django-tree-queries sadly does not (yet?) support ancestors query + # for QS - only for single nodes + + descendants_q = reduce( + operator.or_, + [ + models.Q( + pk__in=entry.ancestors(include_self=include_self).values_list( + "pk", flat=True + ) + ) + for entry in self + ], + initial=models.Q(), + ) + return self.model.objects.filter(descendants_q) + + class Scope(TreeNode, UUIDModel): name = LocalizedCharField(_("scope name"), blank=False, null=False, required=False) @@ -170,7 +219,10 @@ class Scope(TreeNode, UUIDModel): ) is_active = models.BooleanField(default=True) - objects = TreeQuerySet.as_manager(with_tree_fields=True) + objects = ScopeQuerySet.as_manager(with_tree_fields=True) + + def get_root(self): + return self.ancestors(include_self=True).first() def save(self, *args, **kwargs): # django-tree-queries does validation in TreeNode.clean(), which is not