v2 to v3 dump: add associated categories

2 files changed, 24 insertions, 30 deletions

diff --git a/chimere/management/commands/v3_to_v2_markers.py b/chimere/management/commands/v3_to_v2_markers.py
index 0e36370..fd2447b 100644
--- a/chimere/management/commands/v3_to_v2_markers.py
+++ b/chimere/management/commands/v3_to_v2_markers.py
@@ -5,29 +5,28 @@ from django.core.management import call_command
 from django.core.management.base import BaseCommand
 
 
+def get_data(model_name, filters=tuple()):
+    out = StringIO()
+    call_command('dumpdata', 'chimere.' + model_name, stdout=out)
+    values = json.loads(out.getvalue())
+    result = []
+    for value in values:
+        data = value.copy()
+        for attr in filters:
+            data["fields"].pop(attr)
+        result.append(data)
+    return result
+
+
 class Command(BaseCommand):
     help = "Marker dump for v2"
 
     def handle(self, *args, **options):
-        out = StringIO()
-        call_command('dumpdata', 'chimere.Marker', stdout=out)
-        vals = out.getvalue()
-        markers = json.loads(vals)
         new_markers = []
-        for marker in markers:
-            data = marker.copy()
-            data['fields'].pop('weight')
-            data['fields'].pop('normalised_weight')
-            new_markers.append(data)
-        out = StringIO()
-        call_command('dumpdata', 'chimere.Property', stdout=out)
-        vals = out.getvalue()
-        properties = json.loads(vals)
-        for property in properties:
-            data = property.copy()
-            data['fields'].pop('route')
-            data['fields'].pop('polygon')
-            new_markers.append(data)
+        new_markers += get_data('Marker',
+                                filters=('weight', 'normalised_weight'))
+        new_markers += get_data('Marker_categories')
+        new_markers += get_data('Property', filters=('route', 'polygon'))
 
         res = json.dumps(new_markers)
         self.stdout.write(res)
diff --git a/chimere/management/commands/v3_to_v2_routes.py b/chimere/management/commands/v3_to_v2_routes.py
index d373e25..92e26d5 100644
--- a/chimere/management/commands/v3_to_v2_routes.py
+++ b/chimere/management/commands/v3_to_v2_routes.py
@@ -1,25 +1,20 @@
-from io import StringIO
 import json
 
-from django.core.management import call_command
 from django.core.management.base import BaseCommand
 
+from chimere.management.commands.v3_to_v2_markers import get_data
+
 
 class Command(BaseCommand):
     help = "Route dump for v2"
 
     def handle(self, *args, **options):
-        out = StringIO()
-        call_command('dumpdata', 'chimere.Route', indent=4, stdout=out)
-        vals = out.getvalue()
-        routes = json.loads(vals)
         new_routes = []
-        for route in routes:
-            data = route.copy()
-            data['fields'].pop('weight')
-            data['fields'].pop('normalised_weight')
-            data['fields'].pop('color')
-            new_routes.append(data)
+        new_routes += get_data('Route',
+                               filters=('weight', 'normalised_weight',
+                                        'color'))
+        new_routes += get_data('Route_categories')
+
         res = json.dumps(new_routes)
         self.stdout.write(res)