v2 to v3 dump: add associated categories

1 files changed, 17 insertions, 18 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)