Serialization: warehouse - manage filter for context record, site, find and warehouse query

1 files changed, 65 insertions, 62 deletions

diff --git a/archaeological_warehouse/serializers.py b/archaeological_warehouse/serializers.py
index 9a49d5acf..2c44c9dcf 100644
--- a/archaeological_warehouse/serializers.py
+++ b/archaeological_warehouse/serializers.py
@@ -13,6 +13,48 @@ WAREHOUSE_MODEL_LIST = [
 # TODO: associated documents
 
 
+def generate_queryset(base_query_key, ids):
+    warehouse_division_q, warehouse_q, container_q = None, None, None
+    container_loca_q = None
+    for find_key in ("finds", "finds_ref"):
+        base_key = "{}__{}".format(find_key, base_query_key)
+        q = Q(**{base_key: ids})
+        q_loca = Q(
+            **{"container__{}__{}".format(find_key, base_query_key): ids})
+        if not container_q:
+            container_q = q
+            container_loca_q = q_loca
+        else:
+            container_q |= q
+            container_loca_q |= q_loca
+        for container_key in ("containers", "owned_containers"):
+            q = Q(**
+                  {"{}__{}__{}".format(container_key, find_key,
+                                       base_query_key): ids})
+            q_div = Q(**
+                      {"warehouse__{}__{}__{}".format(
+                          container_key, find_key, base_query_key): ids})
+            if not warehouse_q:
+                warehouse_q = q
+                warehouse_division_q = q_div
+            else:
+                warehouse_q |= q
+                warehouse_division_q |= q_div
+    result_queryset = {
+        models.Warehouse.__name__:
+            models.Warehouse.objects.filter(warehouse_q),
+        models.Container.__name__:
+            models.Container.objects.filter(container_q),
+        models.WarehouseDivisionLink.__name__:
+            models.WarehouseDivisionLink.objects.filter(
+                warehouse_division_q),
+        models.ContainerLocalisation.__name__:
+            models.ContainerLocalisation.objects.filter(
+                container_loca_q)
+    }
+    return result_queryset
+
+
 def warehouse_serialization(archive=False, return_empty_types=False,
                             archive_name=None, operation_queryset=None,
                             site_queryset=None, cr_queryset=None,
@@ -21,74 +63,35 @@ def warehouse_serialization(archive=False, return_empty_types=False,
     if operation_queryset:
         operation_ids = operation_queryset.values_list("id", flat=True)
         base_query_key = "base_finds__context_record__operation_id__in"
-        warehouse_division_q, warehouse_q, container_q = None, None, None
-        container_loca_q = None
-        for find_key in ("finds", "finds_ref"):
-            q = Q(**{"{}__{}".format(find_key, base_query_key): operation_ids})
-            q_loca = Q(
-                **{"container__{}__{}".format(find_key, base_query_key):
-                   operation_ids})
-            if not container_q:
-                container_q = q
-                container_loca_q = q_loca
-            else:
-                container_q |= q
-                container_loca_q |= q_loca
-            for container_key in ("containers", "owned_containers"):
-                q = Q(**
-                      {"{}__{}__{}".format(container_key, find_key,
-                                           base_query_key): operation_ids})
-                q_div = Q(**
-                          {"warehouse__{}__{}__{}".format(
-                              container_key, find_key, base_query_key):
-                           operation_ids})
-                if not warehouse_q:
-                    warehouse_q = q
-                    warehouse_division_q = q_div
-                else:
-                    warehouse_q |= q
-                    warehouse_division_q |= q_div
-        result_queryset = {
-            models.Warehouse.__name__:
-                models.Warehouse.objects.filter(warehouse_q),
-            models.Container.__name__:
-                models.Container.objects.filter(container_q),
-            models.WarehouseDivisionLink.__name__:
-                models.WarehouseDivisionLink.objects.filter(
-                    warehouse_division_q),
-            models.ContainerLocalisation.__name__:
-                models.ContainerLocalisation.objects.filter(
-                    container_loca_q)
-        }
+        result_queryset = generate_queryset(base_query_key, operation_ids)
     elif site_queryset:
-        sites = site_queryset.values_list("id", flat=True)
-        f_q = {
-            "base_finds__context_record__operation__archaeological_site_id__in": sites
-        }
-        result_queryset = {
-            models.BaseFind.__name__:
-                models.BaseFind.objects.filter(
-                    context_record__operation__archaeological_site_id__in=sites
-                ),
-            models.Find.__name__:
-                models.Find.objects.filter(**f_q),
-        }
+        site_ids = site_queryset.values_list("id", flat=True)
+        base_query_key = "base_finds__context_record__"\
+                         "archaeological_site_id__in"
+        result_queryset = generate_queryset(base_query_key, site_ids)
     elif cr_queryset:
         cr_ids = cr_queryset.values_list("id", flat=True)
-        result_queryset = {
-            models.BaseFind.__name__:
-                models.BaseFind.objects.filter(
-                    context_record__in=cr_ids),
-            models.Find.__name__:
-                models.Find.objects.filter(
-                    base_finds__context_record__in=cr_ids),
-        }
+        base_query_key = "base_finds__context_record__id__in"
+        result_queryset = generate_queryset(base_query_key, cr_ids)
     elif find_queryset:
         find_ids = find_queryset.values_list("id", flat=True)
+        base_query_key = "id__in"
+        result_queryset = generate_queryset(base_query_key, find_ids)
+    elif warehouse_queryset:
+        warehouse_ids = warehouse_queryset.values_list("id", flat=True)
         result_queryset = {
-            models.BaseFind.__name__: models.BaseFind.objects.filter(
-                find__id__in=find_ids),
-            models.Find.__name__: find_queryset,
+            models.Warehouse.__name__: warehouse_queryset,
+            models.Container.__name__:
+                models.Container.objects.filter(
+                    Q(location__id__in=warehouse_ids) |
+                    Q(responsible__id__in=warehouse_ids)
+                ),
+            models.WarehouseDivisionLink.__name__:
+                models.WarehouseDivisionLink.objects.filter(
+                    warehouse_id__in=warehouse_ids),
+            models.ContainerLocalisation.__name__:
+                models.ContainerLocalisation.objects.filter(
+                    division__warehouse_id__in=warehouse_ids)
         }
     result = generic_get_results(WAREHOUSE_MODEL_LIST, "warehouse",
                                  result_queryset=result_queryset)