Manage multiple condition for cascading shortcut menu - Manage shortcut menu for treatment and treatment files (refs #3384)

1 files changed, 31 insertions, 3 deletions

diff --git a/archaeological_finds/models_treatments.py b/archaeological_finds/models_treatments.py
index eff8d8371..5e763085c 100644
--- a/archaeological_finds/models_treatments.py
+++ b/archaeological_finds/models_treatments.py
@@ -21,7 +21,7 @@ import datetime
 
 from django.conf import settings
 from django.contrib.gis.db import models
-from django.db.models import Max
+from django.db.models import Max, Q
 from django.db.models.signals import post_save, post_delete, pre_delete
 from django.template.defaultfilters import slugify
 from django.utils.translation import ugettext_lazy as _, ugettext
@@ -144,11 +144,29 @@ class Treatment(BaseHistorizedItem, ImageModel, OwnPerms, ShortMenuItem):
             lbl += u" %s %s" % (_(u"by"), unicode(self.person))
         return lbl
 
+    @property
+    def short_class_name(self):
+        return _(u"TREATMENT")
+
+    @classmethod
+    def get_query_owns(cls, user):
+        return (Q(history_creator=user) |
+                Q(person__ishtaruser=user.ishtaruser)) \
+               & Q(end_date__isnull=True)
+
     @classmethod
     def get_owns(cls, user, menu_filtr=None, limit=None):
-        replace_query = {}
+        replace_query = None
         if menu_filtr:
-            replace_query = {'file': menu_filtr}
+            if 'treatmentfile' in menu_filtr:
+                replace_query = Q(file=menu_filtr['treatmentfile'])
+            if 'find' in menu_filtr:
+                q = Q(upstream=menu_filtr['find']) | Q(
+                    downstream=menu_filtr['find'])
+                if replace_query:
+                    replace_query = replace_query | q
+                else:
+                    replace_query = q
         owns = super(Treatment, cls).get_owns(
             user, replace_query=replace_query, limit=limit)
         return sorted(
@@ -455,6 +473,16 @@ class TreatmentFile(ClosedItem, BaseHistorizedItem, OwnPerms, ValueGetter,
         return self.cached_label
 
     @property
+    def short_class_name(self):
+        return _(u"Treatment file")
+
+    @classmethod
+    def get_query_owns(cls, user):
+        return (Q(history_creator=user) |
+                Q(in_charge__ishtaruser=user.ishtaruser)) \
+               & Q(end_date__isnull=True)
+
+    @property
     def associated_filename(self):
         return "-".join([str(slugify(getattr(self, attr)))
                          for attr in ('year', 'index', 'internal_reference',