API Change: All iteration, indexing and slicing operations return Nod…

…e-like objects that have render methods

armstrong/core/arm_wells/models.py

@@ -28,6 +28,12 @@ class Well(models.Model):
 
     objects = WellManager()
 
+    def __init__(self, *args, **kwargs):
+        super(Well, self).__init__(*args, **kwargs)
+        self.queryset = None
+        self.well_content = []
+        self.exclude_ids = []
+
     @property
     def title(self):
         return self.type.title
@@ -48,66 +54,91 @@ def merge_with(self, queryset):
         return self
 
     def render(self, request=None, parent=None):
-        ret = []
-        kwargs = {}
-        if request:
-            kwargs['context_instance'] = RequestContext(request)
-
-        for node in self.nodes.all():
-            kwargs["dictionary"] = {
-                    "well": self,
-                    "object": node.content_object,
-                    "parent": parent,
-            }
-            content = node.content_object
-
-            if hasattr(content, "render"):
-                ret.append(content.render(request, parent=self))
-            else:
-                ret.append(render_to_string("wells/%s/%s/%s.html" % (
-                    content._meta.app_label,
-                    content._meta.object_name.lower(),
-                    self.type.slug), **kwargs))
-        return mark_safe(''.join(ret))
+        if parent is None:
+            parent = self
+        kwargs = {'parent': parent,
+                  'request': request}
+        return mark_safe(''.join(n.render(**kwargs) for n in self))
 
     def initialize(self):
         if self.well_content:
             return
 
-        well_managers = {}
-        for node in self.nodes.all().select_related():
-            model_class = node.content_type.model_class()
-            key = "%s.%s" % (model_class.__module__, node.content_type.model)
-            if not key in well_managers:
-                well_managers[key] = {
-                        "name": node.content_type.model,
-                        "manager": model_class.objects,
-                        "object_ids": [],
-                }
-            self.well_content.append((node.content_type.model, node.object_id))
-            well_managers[key]["object_ids"].append(node.object_id)
-            if self.queryset.model is model_class:
-                self.exclude_ids.append(node.object_id)
-        self.queryset = self.queryset.exclude(pk__in=self.exclude_ids)
+        def process_well(well, managers, ordering):
+            for node in well.nodes.all().select_related():
+                i = len(ordering)
+                model_class = node.content_type.model_class()
+                if model_class == Well:
+                    managers, ordering = process_well(node.content_object,
+                            managers, ordering)
+                    continue
+                key = "%s.%s" % (model_class.__module__, node.content_type.model)
+                if not key in managers:
+                    managers[key] = {
+                            "name": node.content_type.model,
+                            "manager": model_class.objects,
+                            "object_ids": [],
+                    }
+                node_key = "%s.%i" % (node.content_type.model, node.object_id)
+                ordering[node_key] = i, well
+                managers[key]["object_ids"].append(node.object_id)
+                if self.queryset is not None and self.queryset.model is model_class:
+                    self.exclude_ids.append(node.object_id)
+            return managers, ordering
+
+        # well_managers {'module.model':name, manager, ids}
+        self.exclude_ids = []
+        well_managers, ordering = process_well(self, {}, {})
+        self.well_content = [i for i in range(len(ordering))]
 
+        # at this point we know all the queries we need to run to fetch all
+        # content
         for model_data in well_managers.values():
             node_content = model_data["manager"].filter(
                     pk__in=model_data["object_ids"])
             for content in node_content:
-                idx = self.well_content.index((model_data["name"], content.pk))
-                self.well_content[idx] = content
+                node_key = "%s.%i" % (model_data['name'], content.id)
+                idx, well = ordering[node_key]
+                self.well_content[idx] = content, well
+
+        if self.queryset is not None:
+            self.queryset = self.queryset.exclude(pk__in=self.exclude_ids)
+
+
+
+    def __iter__(self):
+        self.initialize()
+        for content, well in self.well_content:
+            yield NodeWrapper(well=well, content_object=content)
+        if self.queryset is not None:
+            for item in self.queryset:
+                yield NodeWrapper(well=self, content_item=item)
 
     def __getslice__(self, i, j):
         self.initialize()
         total_in_well = len(self.well_content)
         if j <= total_in_well:
-            return self.well_content[i:j]
+            return [NodeWrapper(*args) for args in self.well_content[i:j]]
         end = j - total_in_well
         if i >= total_in_well:
             start = i - total_in_well
-            return self.queryset[start:end]
-        return itertools.chain(self.well_content[i:],
-                self.queryset[0:end])
+            return [NodeWrapper(content, self) for content in self.queryset[start:end]]
+        end = j - total_in_well
+        return itertools.chain(
+                (NodeWrapper(*args) for args in self.well_content[i:]),
+                (NodeWrapper(content, self) for content in self.queryset[0:end]))
+
+    def __getitem__(self, i):
+        if type(i) != type(1):
+            raise TypeError
+        self.initialize()
+        con_length = len(self.well_content)
+        if i < con_length:
+            return NodeWrapper(*(self.well_content[i]))
+        elif self.queryset:
+            return NodeWrapper(self.queryset[i-con_length], self)
+        else:
+            raise IndexError
 
     def count(self):
         return self.__len__()
@@ -129,7 +160,25 @@ def __getattr__(self, key):
             raise
 
 
-class Node(models.Model):
+class NodeRenderMixin(object):
+    def render(self, request=None, parent=None):
+        render_args = {
+                'dictionary':{
+                    "well": self.well,
+                    "object": self.content_object,
+                    "parent": parent,
+                },
+            }
+        if parent == self.well:
+            render_args['dictionary']['parent'] = None
+        if request:
+            render_args['context_instance'] = RequestContext(request)
+        return render_to_string("wells/%s/%s/%s.html" % (
+                                self.content_object._meta.app_label,
+                                self.content_object._meta.object_name.lower(),
+                                self.well.type.slug), **render_args)
+
+class Node(NodeRenderMixin, models.Model):
     well = models.ForeignKey(Well, related_name="nodes")
     order = models.IntegerField(default=0)
     content_type = models.ForeignKey(ContentType)
@@ -143,4 +192,11 @@ def __unicode__(self):
         return "%s (%d): %s" % (self.well.title, self.order,
                                 self.content_object)
 
+class NodeWrapper(NodeRenderMixin):
+    def __init__(self, content_object, well):
+        self.well = well
+        self.content_object = content_object
+
+    def __unicode__(self):
+        return "%s: %s" % (self.well.title, self.content_object)
 

armstrong/core/arm_wells/tests/models.py

@@ -207,7 +207,7 @@ def test_well_supports_indexing(self):
         add_n_random_stories_to_well(number_in_well, well)
         i = 0
         for node in well.nodes.all():
-            self.assertEqual(node, well[i])
+            self.assertEqual(node.content_object, well[i].content_object)
             i = i + 1
 
     def test_well_supports_indexing_with_merged_queryset(self):
@@ -221,7 +221,7 @@ def test_well_supports_indexing_with_merged_queryset(self):
         add_n_random_stories_to_well(number_in_well, well)
         i = 0
         for node in well.nodes.all():
-            self.assertEqual(node, well[i])
+            self.assertEqual(node.content_object, well[i].content_object)
             i = i + 1
 
 

armstrong/core/arm_wells/tests/query.py

@@ -6,7 +6,7 @@
         generate_random_image, generate_random_story, generate_random_well,
         TestCase)
 
-from ..models import Node
+from ..models import Node, NodeWrapper
 
 
 class SimpleMergedNodesAndQuerySetTests(TestCase):
@@ -72,16 +72,16 @@ def test_count_and_len_are_identical_with_small_queryset(self):
     def test_node_models_are_first(self):
         node_models = [a.content_object for a in self.well.nodes.all()]
         for obj in self.queryset_backed_well[0:self.number_in_well]:
-            self.assertTrue(obj in node_models)
+            self.assertTrue(obj.content_object in node_models)
 
     def test_fills_in_with_queryset_after_nodes_are_exhausted(self):
         node_models = [a.content_object for a in self.well.nodes.all()]
 
         start = self.number_in_well + 1
         back_filled_models = self.queryset_backed_well[start:]
         for back_filled in back_filled_models:
-            self.assertTrue(isinstance(back_filled, Story), msg="sanity check")
-            self.assertFalse(back_filled in node_models)
+            self.assertTrue(isinstance(back_filled, NodeWrapper), msg="sanity check")
+            self.assertFalse(back_filled.content_object in node_models)
 
     def test_gathers_all_nodes_of_one_type_with_two_queries(self):
         with self.assertNumQueries(2):
@@ -94,9 +94,9 @@ def test_gathers_all_nodes_of_two_types_with_three_queries(self):
 
         number_of_images, number_of_stories = 0, 0
         for node_model in node_models:
-            if node_model.__class__.__name__ == "Story":
+            if node_model.content_object.__class__.__name__ == "Story":
                 number_of_stories += 1
-            if node_model.__class__.__name__ == "Image":
+            if node_model.content_object.__class__.__name__ == "Image":
                 number_of_images += 1
         self.assertEqual(self.number_in_well, number_of_stories)
         self.assertEqual(self.number_in_well, number_of_images)
@@ -113,9 +113,9 @@ def test_perserves_order_across_types(self):
         queryset = well.merge_with(Story.objects.all())
         with self.assertNumQueries(3, msg="sanity check"):
             node_models = queryset[0:3]
-        self.assertEqual(node_models[0], content[0])
-        self.assertEqual(node_models[1], content[1])
-        self.assertEqual(node_models[2], content[2])
+        self.assertEqual(node_models[0].content_object, content[0])
+        self.assertEqual(node_models[1].content_object, content[1])
+        self.assertEqual(node_models[2].content_object, content[2])
 
     def test_does_not_ignore_same_ids_in_different_types(self):
         well = generate_random_well()
@@ -145,15 +145,16 @@ def test_works_with_simple_pagination(self):
         paged = Paginator(self.queryset_backed_well, self.number_in_well)
         page_one = paged.page(1)
         node_models = [a.content_object for a in self.well.nodes.all()]
+        page_one_objects = [n.content_object for n in page_one.object_list]
         for model in node_models:
-            self.assertTrue(model in page_one.object_list)
+            self.assertTrue(model in page_one_objects)
 
     def test_works_when_all_results_are_on_one_page(self):
         paged = Paginator(self.queryset_backed_well, \
                 self.number_in_well + self.number_of_extra_stories)
         page_one = paged.page(1)
         node_models = [a.content_object for a in self.well.nodes.all()]
-        object_list = [a for a in page_one.object_list]
+        object_list = [a.content_object for a in page_one.object_list]
         for model in node_models:
             self.assertTrue(model in object_list)
         for story in self.extra_stories:
@@ -163,5 +164,5 @@ def test_pagination_works_when_split_across_well_and_queryset(self):
         paged = Paginator(self.queryset_backed_well, self.number_in_well + 1)
         page_one = paged.page(1)
         node_models = [a.content_object for a in self.well.nodes.all()]
-        object_list = [a for a in page_one.object_list]
+        object_list = [a.content_object for a in page_one.object_list]
         self.assertFalse(object_list[-1] in node_models)