1 files changed, 57 insertions, 21 deletions

diff --git a/archaeological_operations/models.py b/archaeological_operations/models.py
index cc5c53b45..c6bc25615 100644
--- a/archaeological_operations/models.py
+++ b/archaeological_operations/models.py
@@ -1603,39 +1603,55 @@ class Operation(
 
     DOC_VALUES = [
         ("context_records", _("List of associated context records")),
-        ("containers", _("List of associated containers")),
+        ("containers", _("List of containers with finds related to this operation")),
+        ("document_containers", _("List of containers with documents related to this operation")),
+        ("all_containers", _("List of containers with finds or documents related to this operation"))
     ]
 
-    def get_containers_values(self, filtr, exclude) -> list:  # Container value
-        # list
-        Container = apps.get_model("archaeological_warehouse", "Container")
-        containers = []
-        q = Container.objects.filter(
-            finds__base_finds__context_record__operation=self
-        ).distinct("index")
+    def _get_containers_values(self, query, filtr, exclude) -> list:  # Container value
         exclude += ["operation", "context_record"]
-        for c in q.order_by("index").all():
+        containers = []
+        for c in query.order_by("index").distinct().all():
             containers.append(c.get_values(filtr=filtr, exclude=exclude))
         return containers
 
-    def get_values(self, prefix="", no_values=False, filtr=None, **kwargs):
-        values = super(Operation, self).get_values(
-            prefix=prefix, no_values=no_values, filtr=filtr, **kwargs
-        )
+    def get_containers_values(self, filtr, exclude) -> list:  # Container value
+        return self._get_containers_values(self.containers_q, filtr, exclude)
+
+    def get_document_containers_values(self, filtr, exclude) -> list:  # Container value
+        return self._get_containers_values(self.document_containers_q, filtr, exclude)
+
+    def get_extra_values(self, prefix="", no_values=False, filtr=None, **kwargs):
+        values = {}
         values = get_values_town_related(self, prefix, values, filtr=filtr)
         exclude = kwargs.get("exclude", [])
+        if filtr and "uuid" not in filtr:
+            filtr.append("uuid")
+        containers = []
+        if (not filtr or (f"{prefix}containers" in filtr or f"{prefix}all_containers" in filtr)):
+            containers = self.get_containers_values(filtr, exclude)
+            values[f"{prefix}containers"] = containers[:]
+        document_containers = []
+        if not filtr or (f"{prefix}document_containers" in filtr or f"{prefix}all_containers" in filtr):
+            document_containers = self.get_document_containers_values(filtr, exclude)
+            values[f"{prefix}document_containers"] = document_containers[:]
+        if not filtr or f"{prefix}all_containers" in filtr:
+            values[f"{prefix}all_containers"] = containers[:]
+            current_ids = [container["uuid"] for container in containers]
+            for container in document_containers:
+                if container["uuid"] not in current_ids:
+                    values[f"{prefix}all_containers"].append(container)
         if prefix:
             return values
+        # context_records only when there is no prefix
         if (
-            not filtr or "context_records" in filtr
+                not filtr or "context_records" in filtr
         ) and "context_records" not in exclude:
             kwargs["no_base_finds"] = False
             values["context_records"] = [
                 cr.get_values(prefix=prefix, no_values=True, filtr=None, **kwargs)
                 for cr in self.context_record.all()
             ]
-        if (not filtr or "containers" in filtr) and "context_records" not in exclude:
-            values["containers"] = self.get_containers_values(filtr, exclude)
         return values
 
     def public_representation(self):
@@ -1868,12 +1884,32 @@ class Operation(
             finds__base_finds__context_record__operation=self
         )
 
-    def containers_q(self):
-        from archaeological_warehouse.models import Container
-
-        return Container.objects.filter(
+    @property
+    def containers_q(self):  # -> Container queryset
+        """
+        Containers with finds related to this operation
+        """
+        Container = apps.get_model("archaeological_warehouse", "Container")
+        q = Container.objects.filter(
             finds__base_finds__context_record__operation=self
-        )
+        ).distinct("location", "index")
+        return q
+
+    @property
+    def document_containers_q(self):  # -> Container queryset
+        """
+        Containers with documents related to this operation
+        """
+        q = (Q(operations=self) |
+             Q(context_records__operation=self) |
+             Q(finds__base_finds__context_record__operation=self))
+        q_doc = Document.objects.filter(q)
+
+        Container = apps.get_model("archaeological_warehouse", "Container")
+        q = Container.objects.filter(
+            pk__in=list(q_doc.values_list("container_id", flat=True))
+        ).distinct("location", "index")
+        return q
 
     def get_extra_actions(self, request):
         """