Serialization: warehouse filter

1 files changed, 38 insertions, 1 deletions

diff --git a/archaeological_context_records/serializers.py b/archaeological_context_records/serializers.py
index 7bb9b6724..351b906f4 100644
--- a/archaeological_context_records/serializers.py
+++ b/archaeological_context_records/serializers.py
@@ -1,3 +1,5 @@
+from django.db.models import Q
+
 from ishtar_common.serializers import generic_get_results, archive_serialization
 from archaeological_context_records import models
 
@@ -9,9 +11,41 @@ CR_MODEL_LIST = [
 # TODO: associated documents
 
 
+def generate_warehouse_queryset(ids):
+    base_query_key = "base_finds__find"
+
+    q_cr, q_record_relation = None, None
+
+    for container_key in ("container", "container_ref"):
+        for warehouse_key in ("location", "responsible"):
+            q_c = 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_cr:
+            q_cr = q_c
+            q_record_relation = q_r
+        else:
+            q_cr |= q_c
+            q_record_relation |= q_r
+
+    result_queryset = {
+        models.ContextRecord.__name__: models.ContextRecord.objects.filter(
+            q_cr),
+        models.RecordRelations.__name__:
+            models.RecordRelations.objects.filter(q_r)
+    }
+    return result_queryset
+
+
 def cr_serialization(archive=False, return_empty_types=False,
                      archive_name=None, operation_queryset=None,
-                     site_queryset=None, cr_queryset=None, find_queryset=None):
+                     site_queryset=None, cr_queryset=None, find_queryset=None,
+                     warehouse_queryset=None):
     result_queryset = {}
     if operation_queryset:
         operation_ids = operation_queryset.values_list("id", flat=True)
@@ -59,6 +93,9 @@ def cr_serialization(archive=False, return_empty_types=False,
                     right_record__base_finds__find__id__in=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(CR_MODEL_LIST, "context_records",
                                  result_queryset=result_queryset)