Documentation generation: add some variables to treatments (refs #3586)

2 files changed, 23 insertions, 5 deletions

diff --git a/archaeological_finds/models_finds.py b/archaeological_finds/models_finds.py
index 52601c896..d33933264 100644
--- a/archaeological_finds/models_finds.py
+++ b/archaeological_finds/models_finds.py
@@ -31,7 +31,7 @@ from ishtar_common.utils import cached_label_changed, post_save_point
 
 from ishtar_common.models import GeneralType, ImageModel, BaseHistorizedItem, \
     ShortMenuItem, LightHistorizedItem, HistoricalRecords, OwnPerms, Source, \
-    Person, Basket, get_external_id, post_save_cache
+    Person, Basket, get_external_id, post_save_cache, ValueGetter
 
 from archaeological_operations.models import AdministrativeAct
 from archaeological_context_records.models import ContextRecord, Dating
@@ -490,7 +490,8 @@ class FBulkView(object):
     """
 
 
-class Find(BaseHistorizedItem, ImageModel, OwnPerms, ShortMenuItem):
+class Find(ValueGetter, BaseHistorizedItem, ImageModel, OwnPerms,
+           ShortMenuItem):
     CHECK_DICT = dict(CHECK_CHOICES)
     SHOW_URL = 'show-find'
     SLUG = 'find'
diff --git a/archaeological_finds/models_treatments.py b/archaeological_finds/models_treatments.py
index c128dbe01..0f7c56678 100644
--- a/archaeological_finds/models_treatments.py
+++ b/archaeological_finds/models_treatments.py
@@ -1,6 +1,6 @@
 #!/usr/bin/env python
 # -*- coding: utf-8 -*-
-# Copyright (C) 2016 Étienne Loks  <etienne.loks_AT_peacefrogsDOTnet>
+# Copyright (C) 2016-2017 Étienne Loks  <etienne.loks_AT_peacefrogsDOTnet>
 
 # This program is free software: you can redistribute it and/or modify
 # it under the terms of the GNU Affero General Public License as
@@ -69,7 +69,8 @@ post_save.connect(post_save_cache, sender=TreatmentState)
 post_delete.connect(post_save_cache, sender=TreatmentState)
 
 
-class Treatment(BaseHistorizedItem, ImageModel, OwnPerms, ShortMenuItem):
+class Treatment(ValueGetter, BaseHistorizedItem, ImageModel, OwnPerms,
+                ShortMenuItem):
     SHOW_URL = 'show-treatment'
     TABLE_COLS = ('year', 'index', 'treatment_types__label',
                   'treatment_state__label',
@@ -179,7 +180,8 @@ class Treatment(BaseHistorizedItem, ImageModel, OwnPerms, ShortMenuItem):
         if menu_filtr:
             if 'treatmentfile' in menu_filtr:
                 replace_query = Q(file=menu_filtr['treatmentfile'])
-            if 'find' in menu_filtr and 'basket' not in str(menu_filtr['find']):
+            if 'find' in menu_filtr and \
+                    'basket' not in str(menu_filtr['find']):
                 q = Q(upstream=menu_filtr['find']) | Q(
                     downstream=menu_filtr['find'])
                 if replace_query:
@@ -208,6 +210,21 @@ class Treatment(BaseHistorizedItem, ImageModel, OwnPerms, ShortMenuItem):
         """
         return u" ; ".join([unicode(t) for t in self.treatment_types.all()])
 
+    def get_values(self, prefix=''):
+        values = super(Treatment, self).get_values(prefix=prefix)
+        values[prefix + "upstream_finds"] = u" ; ".join(
+            [unicode(up) for up in self.upstream.all()])
+        values[prefix + "downstream_finds"] = u" ; ".join(
+            [unicode(down) for down in self.downstream.all()])
+        values[prefix + "operations"] = u" ; ".join(
+            [unicode(ope) for ope in self.get_query_operations().all()])
+        if self.upstream.count():
+            find = self.upstream.all()[0]
+            if 'associatedfind_' not in prefix:
+                values.update(
+                    find.get_values(prefix=prefix + 'associatedfind_'))
+        return values
+
     def pre_save(self):
         # is not new
         if self.pk is not None: