Serialization: warehouse filter

7 files changed, 167 insertions, 22 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)
diff --git a/archaeological_finds/models_finds.py b/archaeological_finds/models_finds.py
index 3dda2cdda..6914ff019 100644
--- a/archaeological_finds/models_finds.py
+++ b/archaeological_finds/models_finds.py
@@ -1493,12 +1493,11 @@ class Find(BulkUpdatedItem, ValueGetter, DocumentItem, BaseHistorizedItem,
     datings = models.ManyToManyField(Dating, verbose_name=_(u"Dating"),
                                      related_name='find')
     container = models.ForeignKey(
-        "archaeological_warehouse.Container", verbose_name=_(u"Container"),
-        blank=True, null=True,
-        related_name='finds', on_delete=models.SET_NULL)
+        "archaeological_warehouse.Container", verbose_name=_("Container"),
+        blank=True, null=True, related_name='finds', on_delete=models.SET_NULL)
     container_ref = models.ForeignKey(
         "archaeological_warehouse.Container",
-        verbose_name=_(u"Reference container"),
+        verbose_name=_("Reference container"),
         blank=True, null=True,
         related_name='finds_ref', on_delete=models.SET_NULL)
     is_complete = models.NullBooleanField(_(u"Is complete?"), blank=True,
diff --git a/archaeological_finds/serializers.py b/archaeological_finds/serializers.py
index 7ce17c924..6e079f7ec 100644
--- a/archaeological_finds/serializers.py
+++ b/archaeological_finds/serializers.py
@@ -1,3 +1,5 @@
+from django.db.models import Q
+
 from ishtar_common.serializers import generic_get_results, archive_serialization
 from archaeological_finds import models
 
@@ -9,10 +11,37 @@ FIND_MODEL_LIST = [
 # TODO: associated documents, property, findbasket, treatments
 
 
+def generate_warehouse_queryset(ids):
+    q_find, q_basefind = None, None
+
+    for container_key in ("container", "container_ref"):
+        for warehouse_key in ("location", "responsible"):
+            q_f = Q(**
+                    {"{}__{}__id__in".format(
+                        container_key, warehouse_key): ids})
+            q_bf = Q(**
+                     {"find__{}__{}__id__in".format(
+                         container_key, warehouse_key): ids})
+        if not q_find:
+            q_find = q_f
+            q_basefind = q_bf
+        else:
+            q_find |= q_f
+            q_basefind |= q_bf
+
+    result_queryset = {
+        models.BaseFind.__name__:
+            models.BaseFind.objects.filter(q_basefind),
+        models.Find.__name__:
+            models.Find.objects.filter(q_find),
+    }
+    return result_queryset
+
+
 def find_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)
@@ -54,6 +83,10 @@ def find_serialization(archive=False, return_empty_types=False,
                 find__id__in=find_ids),
             models.Find.__name__: find_queryset,
         }
+    elif warehouse_queryset:
+        warehouse_ids = warehouse_queryset.values_list("id", flat=True)
+        result_queryset = generate_warehouse_queryset(warehouse_ids)
+
     result = generic_get_results(FIND_MODEL_LIST, "finds",
                                  result_queryset=result_queryset)
     full_archive = archive_serialization(
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)
diff --git a/archaeological_warehouse/serializers.py b/archaeological_warehouse/serializers.py
index 2c44c9dcf..ed5a3f189 100644
--- a/archaeological_warehouse/serializers.py
+++ b/archaeological_warehouse/serializers.py
@@ -13,7 +13,7 @@ WAREHOUSE_MODEL_LIST = [
 # TODO: associated documents
 
 
-def generate_queryset(base_query_key, ids):
+def generate_warehouse_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"):
@@ -63,20 +63,21 @@ 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"
-        result_queryset = generate_queryset(base_query_key, operation_ids)
+        result_queryset = generate_warehouse_queryset(base_query_key,
+                                                      operation_ids)
     elif site_queryset:
         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)
+        result_queryset = generate_warehouse_queryset(base_query_key, site_ids)
     elif cr_queryset:
         cr_ids = cr_queryset.values_list("id", flat=True)
         base_query_key = "base_finds__context_record__id__in"
-        result_queryset = generate_queryset(base_query_key, cr_ids)
+        result_queryset = generate_warehouse_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)
+        result_queryset = generate_warehouse_queryset(base_query_key, find_ids)
     elif warehouse_queryset:
         warehouse_ids = warehouse_queryset.values_list("id", flat=True)
         result_queryset = {
diff --git a/archaeological_warehouse/tests.py b/archaeological_warehouse/tests.py
index 52f1a96fc..6d6dbddc9 100644
--- a/archaeological_warehouse/tests.py
+++ b/archaeological_warehouse/tests.py
@@ -32,6 +32,8 @@ from archaeological_finds.models import Find
 from archaeological_warehouse import models, views, forms, serializers
 
 from archaeological_operations.serializers import operation_serialization
+from archaeological_context_records.serializers import cr_serialization
+from archaeological_finds.serializers import find_serialization
 
 
 class SerializationTest(GenericSerializationTest, FindInit, TestCase):
@@ -280,7 +282,6 @@ class SerializationTest(GenericSerializationTest, FindInit, TestCase):
         )
         self.assertEqual(len(loca_json), 2)
 
-    """
     def test_ope_serialization_with_warehouse_filter(self):
         res = self.generic_serialization_test(
             operation_serialization, no_test=True,
@@ -290,18 +291,18 @@ class SerializationTest(GenericSerializationTest, FindInit, TestCase):
         )
         self.assertEqual(len(ope_json), 2)
 
-        result_queryset = models.Find.objects.filter(
-            uuid=self.finds[0].uuid)
+        result_queryset = models.Warehouse.objects.filter(
+            id=self.warehouses[0].pk)
         res = self.generic_serialization_test(
             operation_serialization, no_test=True,
-            kwargs={"find_queryset": result_queryset}
+            kwargs={"warehouse_queryset": result_queryset}
         )
         ope_json = json.loads(
             res[('operations', 'archaeological_operations__Operation')]
         )
         self.assertEqual(len(ope_json), 1)
 
-    def test_cr_serialization_with_find_filter(self):
+    def test_cr_serialization_with_warehouse_filter(self):
         res = self.generic_serialization_test(
             cr_serialization, no_test=True,
         )
@@ -311,18 +312,45 @@ class SerializationTest(GenericSerializationTest, FindInit, TestCase):
         )
         self.assertEqual(len(cr_json), 2)
 
-        result_queryset = models.Find.objects.filter(
-            uuid=self.finds[0].uuid)
+        result_queryset = models.Warehouse.objects.filter(
+            id=self.warehouses[0].pk)
         res = self.generic_serialization_test(
             cr_serialization, no_test=True,
-            kwargs={"find_queryset": result_queryset}
+            kwargs={"warehouse_queryset": result_queryset}
         )
         cr_json = json.loads(
             res[('context_records',
                  'archaeological_context_records__ContextRecord')]
         )
         self.assertEqual(len(cr_json), 1)
-    """
+
+    def test_find_serialization_with_warehouse_filter(self):
+        res = self.generic_serialization_test(
+            find_serialization, no_test=True,
+        )
+        find_json = json.loads(
+            res[('finds', 'archaeological_finds__Find')]
+        )
+        self.assertEqual(len(find_json), 3)
+        bfind_json = json.loads(
+            res[('finds', 'archaeological_finds__BaseFind')]
+        )
+        self.assertEqual(len(bfind_json), 3)
+
+        result_queryset = models.Warehouse.objects.filter(
+            id=self.warehouses[0].pk)
+        res = self.generic_serialization_test(
+            find_serialization, no_test=True,
+            kwargs={"warehouse_queryset": result_queryset}
+        )
+        find_json = json.loads(
+            res[('finds', 'archaeological_finds__Find')]
+        )
+        self.assertEqual(len(find_json), 1)
+        bfind_json = json.loads(
+            res[('finds', 'archaeological_finds__BaseFind')]
+        )
+        self.assertEqual(len(bfind_json), 1)
 
     def test_restore(self):
         current_number, zip_filename = self.generic_restore_test_genzip(
diff --git a/ishtar_common/serializers.py b/ishtar_common/serializers.py
index 36aff680b..034828835 100644
--- a/ishtar_common/serializers.py
+++ b/ishtar_common/serializers.py
@@ -13,9 +13,11 @@ from django.apps import apps
 from django.conf import settings
 from django.contrib.sites.models import Site
 from django.core.serializers import deserialize, serialize
+from django.db.models import Q
 
 from version import get_version
 from . import models
+from archaeological_warehouse import models as warehouse_models
 
 
 class PublicSerializer(serializers.BaseSerializer):
@@ -330,4 +332,4 @@ def restore_serialized(archive_name, delete_existing=False):
                 zip_file.extract("media.zip", tmp_dir_name)
                 with zipfile.ZipFile(
                         tmp_dir_name + os.sep + "media.zip", 'r') as media_zip:
-                    media_zip.extractall(settings.MEDIA_ROOT)
\ No newline at end of file
+                    media_zip.extractall(settings.MEDIA_ROOT)