Serialization: warehouse filter

1 files changed, 46 insertions, 1 deletions

diff --git a/archaeological_operations/serializers.py b/archaeological_operations/serializers.py
index 4dd6f106b..3093bf50a 100644
--- a/archaeological_operations/serializers.py
+++ b/archaeological_operations/serializers.py
@@ -1,3 +1,5 @@
+from django.db.models import Q
+
 from ishtar_common.serializers import generic_get_results, archive_serialization
 from archaeological_operations import models
 
@@ -10,10 +12,50 @@ OPERATION_MODEL_LIST = [
 # TODO: administrativ acts, associated documents
 
 
+def generate_warehouse_queryset(ids):
+    base_query_key = "context_record__base_finds__find"
+
+    q_archaeological_site, q_operation, q_record_relation = None, None, None
+
+    for container_key in ("container", "container_ref"):
+        for warehouse_key in ("location", "responsible"):
+            q_s = Q(**
+                    {"operations__{}__{}__{}__id__in".format(
+                        base_query_key, container_key, warehouse_key): ids})
+            q_o = Q(**
+                    {"{}__{}__{}__id__in".format(
+                        base_query_key, container_key, warehouse_key): ids})
+            q_r = Q(**
+                    {"left_record__{}__{}__{}__id__in".format(
+                        base_query_key, container_key, warehouse_key): ids,
+                     "right_record__{}__{}__{}__id__in".format(
+                         base_query_key, container_key, warehouse_key): ids})
+        if not q_archaeological_site:
+            q_archaeological_site = q_s
+            q_operation = q_o
+            q_record_relation = q_r
+        else:
+            q_archaeological_site |= q_s
+            q_operation |= q_o
+            q_record_relation |= q_r
+
+    result_queryset = {
+        models.ArchaeologicalSite.__name__:
+            models.ArchaeologicalSite.objects.filter(
+                q_archaeological_site
+            ),
+        models.Operation.__name__: models.Operation.objects.filter(
+            q_operation),
+        models.RecordRelations.__name__:
+            models.RecordRelations.objects.filter(q_record_relation),
+    }
+    return result_queryset
+
+
 def operation_serialization(archive=False, return_empty_types=False,
                             archive_name=None, operation_queryset=None,
                             site_queryset=None, cr_queryset=None,
-                            find_queryset=None):
+                            find_queryset=None, warehouse_queryset=None):
     result_queryset = {}
     if operation_queryset:
         operation_ids = operation_queryset.values_list("id", flat=True)
@@ -72,6 +114,9 @@ def operation_serialization(archive=False, return_empty_types=False,
                     find_ids,
                 ),
         }
+    elif warehouse_queryset:
+        warehouse_ids = warehouse_queryset.values_list("id", flat=True)
+        result_queryset = generate_warehouse_queryset(warehouse_ids)
 
     result = generic_get_results(OPERATION_MODEL_LIST, "operations",
                                  result_queryset=result_queryset)