Treatments: manage filter treatment type with hierarchy (do not hide because parent is not OK)

1 files changed, 26 insertions, 0 deletions

diff --git a/archaeological_finds/models_finds.py b/archaeological_finds/models_finds.py
index 8ee742b43..a1321d678 100644
--- a/archaeological_finds/models_finds.py
+++ b/archaeological_finds/models_finds.py
@@ -131,6 +131,32 @@ class TreatmentType(HierarchicalType):
         verbose_name_plural = _("Treatment types")
         ordering = ('order', 'label',)
 
+    @classmethod
+    def get_types(cls, dct=None, instances=False, exclude=None,
+                  empty_first=True, default=None, initial=None, force=False,
+                  full_hierarchy=False):
+        types = super(TreatmentType, cls).get_types(
+            dct=dct, instances=instances, exclude=exclude,
+            empty_first=empty_first, default=default, initial=initial,
+            force=force, full_hierarchy=full_hierarchy)
+        if dct and not exclude:
+            rank = 0
+            if instances:
+                type_list = [ty.pk for ty in types]
+                if types:
+                    rank = types[-1].rank
+            else:
+                type_list = [idx for idx, __ in types]
+            q = cls.objects.filter(**dct).exclude(pk__in=type_list)
+            for t in q.all():
+                if instances:
+                    rank += 1
+                    t.rank = rank
+                    types.append(t)
+                else:
+                    types.append((t.pk, str(t)))
+        return types
+
 
 post_save.connect(post_save_cache, sender=TreatmentType)
 post_delete.connect(post_save_cache, sender=TreatmentType)