Serialization: geo serialization

1 files changed, 54 insertions, 27 deletions

diff --git a/ishtar_common/serializers.py b/ishtar_common/serializers.py
index 4867b58f6..6d25efdb6 100644
--- a/ishtar_common/serializers.py
+++ b/ishtar_common/serializers.py
@@ -103,14 +103,12 @@ def archive_serialization(result, archive_dir=None, archive=False,
     return archive_name
 
 
-def type_serialization(archive=False, return_empty_types=False,
-                       archive_name=None):
+def generic_get_results(model_list, dirname, no_geo=True):
     result = OrderedDict()
-    for model in apps.get_models():
-        if not isinstance(model(), models.GeneralType):
-            continue
+    for model in model_list:
         model_name = model.__name__
         model_name = str(model.__module__).split(".")[0] + "__" + model_name
+
         base_q = model.objects
         q = base_q
         recursion = None
@@ -118,16 +116,25 @@ def type_serialization(archive=False, return_empty_types=False,
             recursion = "parent"
         elif hasattr(model, "inverse_relation"):
             recursion = "inverse_relation"
+        elif hasattr(model, "children"):
+            recursion = "children__id"
         if recursion:
             q = q.filter(**{recursion + "__isnull": True})
-        key = ("types", model_name)
+
+        key = (dirname, model_name)
         result[key] = serialize(
-            "json", q.all(), indent=2,
-            use_natural_foreign_keys=True, use_natural_primary_keys=True
+            "json", q.all(),
+            indent=2,
+            use_natural_foreign_keys=True, use_natural_primary_keys=True,
         )
+
         if recursion:
             serialized = [item["id"] for item in q.values("id").all()]
-            q = base_q.filter(**{recursion + "_id__in": serialized}
+            recursion_in = recursion
+            if not recursion.endswith("_id"):
+                recursion_in += "_id"
+            recursion_in += "__in"
+            q = base_q.filter(**{recursion_in: serialized}
                               ).exclude(id__in=serialized)
             while q.count():
                 v = serialize(
@@ -137,7 +144,7 @@ def type_serialization(archive=False, return_empty_types=False,
                 new_result += json.loads(v)
                 result[key] = json.dumps(new_result, indent=2)
                 serialized += [item["id"] for item in q.values("id").all()]
-                q = base_q.filter(**{recursion + "_id__in": serialized}
+                q = base_q.filter(**{recursion_in: serialized}
                                   ).exclude(id__in=serialized)
             # managed circular
             q = base_q.exclude(id__in=serialized)
@@ -153,27 +160,20 @@ def type_serialization(archive=False, return_empty_types=False,
                 new_result += result_cleaned
                 new_result += result_to_add
                 result[key] = json.dumps(new_result, indent=2)
-    return archive_serialization(result, archive_dir="types", archive=archive,
-                                 return_empty_types=return_empty_types,
-                                 archive_name=archive_name)
-
 
-def generic_get_results(model_list, dirname):
-    result = OrderedDict()
-    for model in model_list:
-        model_name = model.__name__
-        model_name = str(model.__module__).split(".")[0] + "__" + model_name
-        key = (dirname, model_name)
-        result[key] = serialize(
-            "json", model.objects.all(),
-            indent=2,
-            use_natural_foreign_keys=True, use_natural_primary_keys=True,
-        )
+        excluded_fields = []
         if hasattr(model, "SERIALIZATION_EXCLUDE"):
+            excluded_fields = list(model.SERIALIZATION_EXCLUDE)
+        if no_geo:
+            excluded_fields += ["center", "limit"] + [
+                field.name for field in models.GeoItem._meta.get_fields()
+            ]
+        if excluded_fields:
             new_result = json.loads(result[key])
             for idx in range(len(new_result)):
-                for excluded_field in model.SERIALIZATION_EXCLUDE:
-                    new_result[idx]["fields"].pop(excluded_field)
+                for excluded_field in excluded_fields:
+                    if excluded_field in new_result[idx]["fields"]:
+                        new_result[idx]["fields"].pop(excluded_field)
             result[key] = json.dumps(new_result, indent=2)
     return result
 
@@ -201,6 +201,18 @@ def generic_archive_files(model_list, archive_name=None):
     return archive_name
 
 
+def type_serialization(archive=False, return_empty_types=False,
+                       archive_name=None):
+    TYPE_MODEL_LIST = [
+        model for model in apps.get_models()
+        if isinstance(model(), models.GeneralType)
+    ]
+    result = generic_get_results(TYPE_MODEL_LIST, "types")
+    return archive_serialization(result, archive_dir="types", archive=archive,
+                                 return_empty_types=return_empty_types,
+                                 archive_name=archive_name)
+
+
 CONF_MODEL_LIST = [
     models.IshtarSiteProfile, models.GlobalVar, models.CustomForm,
     models.ExcludedField, models.JsonDataSection, models.JsonDataField,
@@ -243,6 +255,20 @@ def importer_serialization(archive=False, return_empty_types=False,
     return full_archive
 
 
+GEO_MODEL_LIST = [
+    models.State, models.Department, models.Town
+]
+
+
+def geo_serialization(archive=False, return_empty_types=False,
+                      archive_name=None, no_geo=True):
+    result = generic_get_results(GEO_MODEL_LIST, "common_geo", no_geo=no_geo)
+    full_archive = archive_serialization(
+        result, archive_dir="common_geo", archive=archive,
+        return_empty_types=return_empty_types, archive_name=archive_name)
+    return full_archive
+
+
 def restore_serialized(archive_name, delete_existing=False):
     with zipfile.ZipFile(archive_name, "r") as zip_file:
         # check version
@@ -256,6 +282,7 @@ def restore_serialized(archive_name, delete_existing=False):
         DIRS = (
             ("types", [None]), ("common_configuration", CONF_MODEL_LIST),
             ("common_imports", IMPORT_MODEL_LIST),
+            ("common_geo", GEO_MODEL_LIST)
         )
         namelist = zip_file.namelist()
         for current_dir, model_list in DIRS: