Simple management for shortcut menu of treatment and treatment file (refs #3384)

4 files changed, 39 insertions, 6 deletions

diff --git a/archaeological_finds/models_treatments.py b/archaeological_finds/models_treatments.py
index 67bbd3ddb..eff8d8371 100644
--- a/archaeological_finds/models_treatments.py
+++ b/archaeological_finds/models_treatments.py
@@ -30,7 +30,7 @@ from django.utils.translation import ugettext_lazy as _, ugettext
 from ishtar_common.utils import cached_label_changed
 from ishtar_common.models import GeneralType, ImageModel, BaseHistorizedItem, \
     OwnPerms, HistoricalRecords, Person, Organization, Source, \
-    ValueGetter, post_save_cache
+    ValueGetter, post_save_cache, ShortMenuItem
 from archaeological_warehouse.models import Warehouse, Container
 from archaeological_finds.models_finds import Find, FindBasket
 from archaeological_operations.models import ClosedItem
@@ -66,7 +66,7 @@ post_save.connect(post_save_cache, sender=TreatmentState)
 post_delete.connect(post_save_cache, sender=TreatmentState)
 
 
-class Treatment(BaseHistorizedItem, ImageModel, OwnPerms):
+class Treatment(BaseHistorizedItem, ImageModel, OwnPerms, ShortMenuItem):
     SHOW_URL = 'show-treatment'
     TABLE_COLS = ('year', 'index', 'treatment_types__label',
                   'treatment_state__label',
@@ -87,6 +87,7 @@ class Treatment(BaseHistorizedItem, ImageModel, OwnPerms):
         "upstream_cached_label": _(u"Upstream find"),
     }
     IMAGE_PREFIX = 'treatment'
+    SLUG = 'treatment'
     label = models.CharField(_(u"Label"), blank=True, null=True,
                              max_length=200)
     other_reference = models.CharField(_(u"Other ref."), blank=True, null=True,
@@ -143,6 +144,17 @@ class Treatment(BaseHistorizedItem, ImageModel, OwnPerms):
             lbl += u" %s %s" % (_(u"by"), unicode(self.person))
         return lbl
 
+    @classmethod
+    def get_owns(cls, user, menu_filtr=None, limit=None):
+        replace_query = {}
+        if menu_filtr:
+            replace_query = {'file': menu_filtr}
+        owns = super(Treatment, cls).get_owns(
+            user, replace_query=replace_query, limit=limit)
+        return sorted(
+            owns, key=lambda x: x.cached_label
+            if hasattr(x, 'cached_label') else unicode(x))
+
     def _generate_cached_label(self):
         items = [unicode(getattr(self, k))
                  for k in ['year', 'index', 'other_reference', 'label'] if
@@ -382,10 +394,12 @@ post_save.connect(post_save_cache, sender=TreatmentFileType)
 post_delete.connect(post_save_cache, sender=TreatmentFileType)
 
 
-class TreatmentFile(ClosedItem, BaseHistorizedItem, OwnPerms, ValueGetter):
+class TreatmentFile(ClosedItem, BaseHistorizedItem, OwnPerms, ValueGetter,
+                      ShortMenuItem):
     SLUG = 'treatmentfile'
     SHOW_URL = 'show-treatmentfile'
     TABLE_COLS = ['type', 'year', 'index', 'internal_reference', 'name']
+    SLUG = 'treatmentfile'
 
     # fields
     year = models.IntegerField(_(u"Year"),
@@ -446,6 +460,13 @@ class TreatmentFile(ClosedItem, BaseHistorizedItem, OwnPerms, ValueGetter):
                          for attr in ('year', 'index', 'internal_reference',
                                       'name') if getattr(self, attr)])
 
+    @classmethod
+    def get_owns(cls, user, menu_filtr=None, limit=None):
+        owns = super(TreatmentFile, cls).get_owns(user, limit=limit)
+        return sorted(
+            owns, key=lambda x: x.cached_label
+            if hasattr(x, 'cached_label') else unicode(x))
+
     def _generate_cached_label(self):
         items = [unicode(getattr(self, k))
                  for k in ['year', 'index', 'internal_reference', 'name'] if
diff --git a/archaeological_finds/urls.py b/archaeological_finds/urls.py
index f1ddb1223..7be07d015 100644
--- a/archaeological_finds/urls.py
+++ b/archaeological_finds/urls.py
@@ -275,10 +275,16 @@ urlpatterns += patterns(
         'revert_find', name='revert-find'),
     url(r'^get-treatment/(?P<type>.+)?$',
         'get_treatment', name='get-treatment'),
+    url(r'get-treatment-shortcut/(?P<type>.+)?$',
+        'get_treatment', name='get-treatment-shortcut',
+        kwargs={'full': 'shortcut'}),
     url(r'^show-treatment(?:/(?P<pk>.+))?/(?P<type>.+)?$', 'show_treatment',
         name=models.Treatment.SHOW_URL),
     url(r'get-treatmentfile/(?P<type>.+)?$',
         'get_treatmentfile', name='get-treatmentfile'),
+    url(r'get-treatmentfile-shortcut/(?P<type>.+)?$',
+        'get_treatmentfile', name='get-treatmentfile-shortcut',
+        kwargs={'full': 'shortcut'}),
     url(r'^show-treatmentfile(?:/(?P<pk>.+))?/(?P<type>.+)?$',
         'show_treatmentfile',
         name=models.TreatmentFile.SHOW_URL),
diff --git a/ishtar_common/models.py b/ishtar_common/models.py
index d05ffdd77..fcb367ae5 100644
--- a/ishtar_common/models.py
+++ b/ishtar_common/models.py
@@ -254,7 +254,8 @@ class OwnPerms:
         query = cls.get_query_owns(user)
         if not query and not replace_query:
             return cls.objects.filter(pk__isnull=True)
-        q = cls.objects.filter(query)
+        if query:
+            q = cls.objects.filter(query)
         if replace_query:
             q = cls.objects.filter(**replace_query)
         if limit:
diff --git a/ishtar_common/views.py b/ishtar_common/views.py
index 6426fef8f..783c6badb 100644
--- a/ishtar_common/views.py
+++ b/ishtar_common/views.py
@@ -58,7 +58,7 @@ from menus import menu
 from archaeological_files.models import File
 from archaeological_operations.models import Operation
 from archaeological_context_records.models import ContextRecord
-from archaeological_finds.models import Find
+from archaeological_finds.models import Find, Treatment, TreatmentFile
 
 from archaeological_operations.forms import DashboardForm as DashboardFormOpe
 from archaeological_files.forms import DashboardForm as DashboardFormFile
@@ -239,6 +239,9 @@ def shortcut_menu(request):
         CURRENT_ITEMS.append((_(u"Context record"), ContextRecord))
     if profile.find:
         CURRENT_ITEMS.append((_(u"Find"), Find))
+    if profile.warehouse:
+        CURRENT_ITEMS.append((_(u"Treatment file"), TreatmentFile))
+        CURRENT_ITEMS.append((_(u"Treatment"), Treatment))
     if hasattr(request.user, 'ishtaruser') and \
             request.user.ishtaruser.advanced_shortcut_menu:
         dct = {'current_menu': [], 'menu': [],
@@ -301,7 +304,9 @@ def get_current_items(request):
     for key, model in (('file', File),
                        ('operation', Operation),
                        ('contextrecord', ContextRecord),
-                       ('find', Find)):
+                       ('find', Find),
+                       ('treatmentfile', TreatmentFile),
+                       ('treatment', Treatment)):
         currents[key] = None
         if key in request.session and request.session[key]:
             try: