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

2 files changed, 30 insertions, 3 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),