Merge pull request #2 from solarmonkey/fix-nested-dumps

Fix nested dumps

fixture_magic/management/commands/custom_dump.py

@@ -1,7 +1,5 @@
 from __future__ import print_function
 
-import sys
-
 try:
     import json
 except ImportError:
@@ -14,16 +12,35 @@
     from django.apps import apps as loading
 from django.core.serializers import serialize
 from django.conf import settings
-from django.template import Variable, VariableDoesNotExist
 
 from fixture_magic.utils import (
     add_to_serialize_list,
     reorder_json,
-    serialize_me,
     serialize_fully
 )
 
 
+def process_dep(parent, dep, serialize_me, seen):
+    parts = dep.split('.')
+    current = parts.pop(0)
+    leftover = '.'.join(parts)
+
+    try:
+        thing = getattr(parent, current)
+    except AttributeError:
+        pass  # related object not found
+    else:
+        if hasattr(thing, 'all'):
+            children = thing.all()
+        else:
+            children = [thing]
+        add_to_serialize_list(children, serialize_me, seen)
+
+        if leftover:
+            for child in children:
+                process_dep(child, leftover, serialize_me, seen)
+
+
 class Command(BaseCommand):
     help = 'Dump multiple pre-defined sets of objects into a JSON fixture.'
 
@@ -34,46 +51,43 @@ def add_arguments(self, parser):
                             help='Use natural keys if they are available.')
 
     def handle(self, *args, **options):
+        serialize_me = []
+        seen = set()
         # Get the primary object
         dump_name = options['dump_name']
         pks = options['pk']
         dump_settings = settings.CUSTOM_DUMPS[dump_name]
         app_label, model_name, *manager_method = dump_settings['primary'].split('.')
         include_primary = dump_settings.get("include_primary", False)
+
         default_manager = loading.get_model(app_label, model_name).objects
         if manager_method:
-            dump_me = getattr(default_manager, manager_method[0])()
+            queryset = getattr(default_manager, manager_method[0])()
         else:
-            dump_me = default_manager
+            queryset = default_manager.all()
         if pks:
-            objs = dump_me.filter(pk__in=pks)
-        else:
-            objs = dump_me
-        for obj in objs.all():
+            queryset = queryset.filter(pk__in=pks)
+
+        deps = dump_settings.get('dependents', [])
+        for obj in queryset:
             # get the dependent objects and add to serialize list
-            for dep in dump_settings['dependents']:
-                try:
-                    thing = Variable("thing.%s" % dep).resolve({'thing': obj})
-                    if hasattr(thing, 'all'):  # Related managers can't be iterated over
-                        thing = thing.all()
-                    add_to_serialize_list([thing])
-                except VariableDoesNotExist:
-                    sys.stderr.write('%s not found' % dep)
-
-            if include_primary or not dump_settings['dependents']:
-                add_to_serialize_list([obj])
-
-        serialize_fully()
-        data = serialize('json', [o for o in serialize_me if o is not None],
-                         indent=4,
-                         use_natural_foreign_keys=options.get('natural', False),
-                         use_natural_primary_keys=options.get('natural', False),
-                         )
+            for dep in deps:
+                process_dep(obj, dep, serialize_me, seen)
+            if include_primary or not deps:
+                add_to_serialize_list([obj], serialize_me, seen)
+
+        serialize_fully(serialize_me, seen)
+        data = serialize(
+            'json', [o for o in serialize_me if o is not None],
+            indent=4,
+            use_natural_foreign_keys=options.get('natural', False),
+            use_natural_primary_keys=options.get('natural', False),
+        )
 
         data = reorder_json(
             json.loads(data),
             dump_settings.get('order', []),
             ordering_cond=dump_settings.get('order_cond', {})
         )
 
-        print(json.dumps(data, indent=4))
+        self.stdout.write(json.dumps(data, indent=4))

fixture_magic/management/commands/dump_object.py

@@ -12,8 +12,7 @@
     from django.apps import apps as loading
 import json
 
-from fixture_magic.utils import (add_to_serialize_list, serialize_me, seen,
-                                 serialize_fully)
+from fixture_magic.utils import (add_to_serialize_list, serialize_fully)
 
 
 class Command(BaseCommand):
@@ -65,6 +64,8 @@ def add_arguments(self, parser):
         )
 
     def handle(self, *args, **options):
+        serialize_me = []
+        seen = set()
         error_text = ('%s\nTry calling dump_object with --help argument or ' +
                       'use the following arguments:\n %s' % self.args)
         try:
@@ -117,18 +118,18 @@ def handle(self, *args, **options):
                 for rel in related_fields:
                     try:
                         if hasattr(getattr(obj, rel), 'all'):
-                            add_to_serialize_list(getattr(obj, rel).all())
+                            add_to_serialize_list(getattr(obj, rel).all(), serialize_me, seen)
                         else:
-                            add_to_serialize_list([getattr(obj, rel)])
+                            add_to_serialize_list([getattr(obj, rel)], serialize_me, seen)
                     except FieldError:
                         pass
                     except ObjectDoesNotExist:
                         pass
 
-        add_to_serialize_list(objs)
+        add_to_serialize_list(objs, serialize_me, seen)
 
         if options.get('follow_fk', True):
-            serialize_fully()
+            serialize_fully(serialize_me, seen)
         else:
             # reverse list to match output of serializez_fully
             serialize_me.reverse()
@@ -143,8 +144,3 @@ def handle(self, *args, **options):
                                     indent=4,
                                     use_natural_foreign_keys=natural_foreign,
                                     use_natural_primary_keys=natural_primary))
-
-        # Clear the list. Useful for when calling multiple
-        # dump_object commands with a single execution of django
-        del serialize_me[:]
-        seen.clear()

fixture_magic/utils.py

@@ -1,8 +1,5 @@
 from django.db import models
 
-serialize_me = []
-seen = {}
-
 
 def reorder_json(data, models, ordering_cond=None):
     """Reorders JSON (actually a list of model dicts).
@@ -45,21 +42,24 @@ def get_fields(obj):
         return []
 
 
-def serialize_fully():
+def serialize_fully(serialize_me, seen):
     index = 0
 
     while index < len(serialize_me):
         for field in get_fields(serialize_me[index]):
             if isinstance(field, models.ForeignKey):
                 add_to_serialize_list(
-                    [serialize_me[index].__getattribute__(field.name)])
+                    [serialize_me[index].__getattribute__(field.name)],
+                    serialize_me,
+                    seen
+                )
 
         index += 1
 
     serialize_me.reverse()
 
 
-def add_to_serialize_list(objs):
+def add_to_serialize_list(objs, serialize_me, seen):
     for obj in objs:
         if obj is None:
             continue
@@ -76,4 +76,4 @@ def add_to_serialize_list(objs):
 
         if key not in seen:
             serialize_me.append(obj)
-            seen[key] = 1
+            seen.add(key)