Remove recursion from BaseNestedSets.get_tree method #39

Makefile

@@ -1,4 +1,7 @@
 all: test
 
 test:
-	nosetests --with-coverage --nocapture --cover-package=sqlalchemy_mptt --cover-erase --with-doctest
+	nosetests --with-coverage --cover-package=sqlalchemy_mptt --cover-erase --with-doctest
+
+nocapture:
+	nosetests --with-coverage --cover-package=sqlalchemy_mptt --cover-erase --with-doctest --nocapture

sqlalchemy_mptt/mixins.py

@@ -88,12 +88,13 @@ def parent_id(cls):
     @declared_attr
     def parent(cls):
         pk = getattr(cls, cls.get_pk())
-        return relationship(cls, primaryjoin=lambda: pk == cls.parent_id,
-                            order_by=lambda: cls.left,
-                            backref=backref('children', cascade="all,delete",
-                                            order_by=lambda: cls.left),
-                            remote_side=cls.get_class_pk(),  # for show in sacrud relation
-                            )
+        return relationship(
+            cls, primaryjoin=lambda: pk == cls.parent_id,
+            order_by=lambda: cls.left,
+            backref=backref('children', cascade="all,delete",
+                            order_by=lambda: cls.left),
+            remote_side=cls.get_class_pk(),  # for show in sacrud relation
+        )
 
     @declared_attr
     def left(cls):
@@ -133,7 +134,7 @@ def move_inside(self, parent_id):
 
         * :mod:`sqlalchemy_mptt.tests.TestTree.test_move_inside_function`
         * :mod:`sqlalchemy_mptt.tests.TestTree.test_move_inside_to_the_same_parent_function`
-        """
+        """  # noqa
         session = Session.object_session(self)
         self.parent_id = parent_id
         self.mptt_move_inside = parent_id
@@ -143,7 +144,7 @@ def move_after(self, node_id):
         """ Moving one node of tree after another
 
         For example see :mod:`sqlalchemy_mptt.tests.TestTree.test_move_after_function`
-        """
+        """  # noqa
         session = Session.object_session(self)
         self.parent_id = self.parent_id
         self.mptt_move_after = node_id
@@ -170,7 +171,7 @@ def leftsibling_in_level(self):
         """ Node to the left of the current node at the same level
 
         For example see :mod:`sqlalchemy_mptt.tests.TestTree.test_leftsibling_in_level`
-        """
+        """  # noqa
         table = _get_tree_table(self.__mapper__)
         session = Session.object_session(self)
         current_lvl_nodes = session.query(table)\
@@ -181,8 +182,23 @@ def leftsibling_in_level(self):
         return None
 
     @classmethod
-    def get_tree(cls, session, json=False, json_fields=None):
-        """ This function generate tree of current node in dict or json format.
+    def _get_tree_node(cls, node, json, json_fields):
+        """ Helper method for ``get_tree`` and ``get_tree_reqursively``.
+        """
+        if json:
+            pk = getattr(node, node.get_pk())
+            # jqTree or jsTree format
+            result = {'id': pk, 'label': node.__repr__()}
+            if json_fields:
+                result.update(json_fields(node))
+        else:
+            result = {'node': node}
+        return result
+
+    @classmethod
+    def get_tree_reqursively(cls, session, json=False, json_fields=None):
+        """ This function recursively generate tree of current node in dict or
+        json format.
 
         Args:
             session (:mod:`sqlalchemy.orm.session.Session`): SQLAlchemy session
@@ -196,15 +212,9 @@ def get_tree(cls, session, json=False, json_fields=None):
         * :mod:`sqlalchemy_mptt.tests.TestTree.test_get_tree`
         * :mod:`sqlalchemy_mptt.tests.TestTree.test_get_json_tree`
         * :mod:`sqlalchemy_mptt.tests.TestTree.test_get_json_tree_with_custom_field`
-        """
+        """  # noqa
         def recursive_node_to_dict(node):
-            result = {'node': node}
-            pk = getattr(node, node.get_pk())
-            if json:
-                # jqTree or jsTree format
-                result = {'id': pk, 'label': node.__repr__()}
-                if json_fields:
-                    result.update(json_fields(node))
+            result = cls._get_tree_node(node, json, json_fields)
             children = [recursive_node_to_dict(c) for c in node.children]
             if children:
                 result['children'] = children
@@ -218,6 +228,48 @@ def recursive_node_to_dict(node):
 
         return tree
 
+    @classmethod
+    def get_tree(cls, session, json=False, json_fields=None):
+        """ This function generate tree of current node in dict or json format.
+
+        Args:
+            session (:mod:`sqlalchemy.orm.session.Session`): SQLAlchemy session
+
+        Kwargs:
+            json (bool): if True return JSON jqTree format
+            json_fields (function): append custom fields in JSON
+
+        Example:
+
+        * :mod:`sqlalchemy_mptt.tests.TestTree.test_get_tree`
+        * :mod:`sqlalchemy_mptt.tests.TestTree.test_get_json_tree`
+        * :mod:`sqlalchemy_mptt.tests.TestTree.test_get_json_tree_with_custom_field`
+        """  # noqa
+        nodes = session.query(cls).order_by(cls.level).all()
+        tree = []
+        nodes_of_level = {}
+
+        def get_node_id(node):
+            return getattr(node, node.get_pk())
+
+        for node in nodes:
+            result = cls._get_tree_node(node, json, json_fields)
+            parent_id = node.parent_id
+            # Parent detect!
+            if parent_id:
+                # Find parent in tree list!
+                if parent_id in nodes_of_level.keys():
+                    if 'children' not in nodes_of_level[parent_id]:
+                        nodes_of_level[parent_id]['children'] = []
+                    # Append to parent!
+                    nl = nodes_of_level[parent_id]['children']
+                    nl.append(result)
+                    nodes_of_level[get_node_id(node)] = nl[-1]
+            else:
+                tree.append(result)
+                nodes_of_level[get_node_id(node)] = tree[-1]
+        return tree
+
     @classmethod
     def rebuild_tree(cls, session, tree_id):
         """ This function rebuid tree.

sqlalchemy_mptt/tests/tree_testing_base.py

@@ -7,7 +7,7 @@
 # Distributed under terms of the MIT license.
 
 
-from sqlalchemy import create_engine
+from sqlalchemy import create_engine, event
 from sqlalchemy.exc import IntegrityError
 from sqlalchemy.orm import sessionmaker
 
@@ -103,6 +103,18 @@ class TreeTestingMixin(object):
     base = None
     model = None
 
+    def catch_queries(self, conn, cursor, statement, *args):
+        self.stmts.append(statement)
+
+    def start_query_counter(self):
+        self.stmts = []
+        event.listen(self.session.bind.engine, "before_cursor_execute",
+                     self.catch_queries)
+
+    def stop_query_counter(self):
+        event.remove(self.session.bind.engine, "before_cursor_execute",
+                     self.catch_queries)
+
     def setUp(self):
         self.engine = create_engine('sqlite:///:memory:')
         Session = mptt_sessionmaker(sessionmaker(bind=self.engine))
@@ -754,7 +766,7 @@ def test_move_tree_to_another_tree(self):
                                                                                      |         |
             6                                                                    26(20)27  30(22)31
 
-        """
+        """  # noqa
         node = self.session.query(self.model).\
             filter(self.model.ppk == 12).one()
         node.parent_id = 7
@@ -928,7 +940,8 @@ def test_move_to_toplevel_where_much_trees_from_right_side(self):
             4               8(20)9    12(22)13
 
         """
-        node = self.session.query(self.model).filter(self.model.ppk == 15).one()
+        node = self.session.query(self.model)\
+            .filter(self.model.ppk == 15).one()
         node.move_after("1")
         #                 id lft rgt lvl parent tree
         self.assertEqual([(1,   1, 22, 1, None, 1),
@@ -957,7 +970,8 @@ def test_move_to_toplevel_where_much_trees_from_right_side(self):
                           (21, 11, 14, 3, 18, 3),
                           (22, 12, 13, 4, 21, 3)], self.result.all())
 
-        node = self.session.query(self.model).filter(self.model.ppk == 20).one()
+        node = self.session.query(self.model)\
+            .filter(self.model.ppk == 20).one()
         node.move_after("1")
         """ level           tree_id = 1
             1                    1(1)22
@@ -1254,11 +1268,55 @@ def test_get_tree(self):
             tree = Tree.get_tree(self.session)
         """
         tree = self.model.get_tree(self.session)
+        tree_reqursively = self.model.get_tree_reqursively(self.session)
 
         def go(id):
             return get_obj(self.session, self.model, id)
-        self.assertEqual(tree,
-                         [{'node': go(1), 'children': [{'node': go(2), 'children': [{'node': go(3)}]}, {'node': go(4), 'children': [{'node': go(5)}, {'node': go(6)}]}, {'node': go(7), 'children': [{'node': go(8), 'children': [{'node': go(9)}]}, {'node': go(10), 'children': [{'node': go(11)}]}]}]}, {'node': go(12), 'children': [{'node': go(13), 'children': [{'node': go(14)}]}, {'node': go(15), 'children': [{'node': go(16)}, {'node': go(17)}]}, {'node': go(18), 'children': [{'node': go(19), 'children': [{'node': go(20)}]}, {'node': go(21), 'children': [{'node': go(22)}]}]}]}])
+
+        reference_tree = [{'node': go(1), 'children': [{'node': go(2), 'children': [{'node': go(3)}]}, {'node': go(4), 'children': [{'node': go(5)}, {'node': go(6)}]}, {'node': go(7), 'children': [{'node': go(8), 'children': [{'node': go(9)}]}, {'node': go(10), 'children': [{'node': go(11)}]}]}]}, {'node': go(12), 'children': [{'node': go(13), 'children': [{'node': go(14)}]}, {'node': go(15), 'children': [{'node': go(16)}, {'node': go(17)}]}, {'node': go(18), 'children': [{'node': go(19), 'children': [{'node': go(20)}]}, {'node': go(21), 'children': [{'node': go(22)}]}]}]}]  # noqa
+
+        self.assertEqual(tree, reference_tree)
+        self.assertEqual(tree_reqursively, reference_tree)
+
+    def test_get_tree_count_query(self):
+        """
+        Count num of queries to the database.
+        See https://github.com/ITCase/sqlalchemy_mptt/issues/39
+
+
+        Use ``--nocapture`` option for show run time:
+
+        ::
+
+            nosetests sqlalchemy_mptt.tests.test_events:TestTree.test_get_tree_count_query --nocapture
+            Get tree:             0:00:00.001817
+            Get tree reqursively: 0:00:00.020615
+            .
+            ----------------------------------------------------------------------
+            Ran 1 test in 0.064s
+
+            OK
+        """  # noqa
+        from datetime import datetime
+        self.session.commit()
+
+        # Get tree by for cycle
+        self.start_query_counter()
+        self.assertEqual(0, len(self.stmts))
+        startTime = datetime.now()
+        self.model.get_tree(self.session)
+        print("Get tree: {!s:>26}".format(datetime.now() - startTime))
+        self.assertEqual(1, len(self.stmts))
+        self.stop_query_counter()
+
+        # Get tree by recursion
+        self.start_query_counter()
+        self.assertEqual(0, len(self.stmts))
+        startTime = datetime.now()
+        self.model.get_tree_reqursively(self.session)
+        print("Get tree reqursively: {}".format(datetime.now() - startTime))
+        self.assertEqual(23, len(self.stmts))
+        self.stop_query_counter()
 
     def test_get_json_tree(self):
         """.. note::
@@ -1271,9 +1329,13 @@ def test_get_json_tree(self):
 
             tree = Tree.get_tree(self.session, json=True)
         """
+        reference_tree = [{'children': [{'children': [{'id': 3, 'label': '<Node (3)>'}], 'id': 2, 'label': '<Node (2)>'}, {'children': [{'id': 5, 'label': '<Node (5)>'}, {'id': 6, 'label': '<Node (6)>'}], 'id': 4, 'label': '<Node (4)>'}, {'children': [{'children': [{'id': 9, 'label': '<Node (9)>'}], 'id': 8, 'label': '<Node (8)>'}, {'children': [{'id': 11, 'label': '<Node (11)>'}], 'id': 10, 'label': '<Node (10)>'}], 'id': 7, 'label': '<Node (7)>'}], 'id': 1, 'label': '<Node (1)>'}, {'children': [{'children': [{'id': 14, 'label': '<Node (14)>'}], 'id': 13, 'label': '<Node (13)>'}, {'children': [{'id': 16, 'label': '<Node (16)>'}, {'id': 17, 'label': '<Node (17)>'}], 'id': 15, 'label': '<Node (15)>'}, {'children': [{'children': [{'id': 20, 'label': '<Node (20)>'}], 'id': 19, 'label': '<Node (19)>'}, {'children': [{'id': 22, 'label': '<Node (22)>'}], 'id': 21, 'label': '<Node (21)>'}], 'id': 18, 'label': '<Node (18)>'}], 'id': 12, 'label': '<Node (12)>'}]  # noqa
+
         tree = self.model.get_tree(self.session, json=True)
-        self.assertEqual(tree, [{'children': [{'children': [{'id': 3, 'label': '<Node (3)>'}], 'id': 2, 'label': '<Node (2)>'}, {'children': [{'id': 5, 'label': '<Node (5)>'}, {'id': 6, 'label': '<Node (6)>'}], 'id': 4, 'label': '<Node (4)>'}, {'children': [{'children': [{'id': 9, 'label': '<Node (9)>'}], 'id': 8, 'label': '<Node (8)>'}, {'children': [{'id': 11, 'label': '<Node (11)>'}], 'id': 10, 'label': '<Node (10)>'}], 'id': 7, 'label': '<Node (7)>'}], 'id': 1, 'label': '<Node (1)>'}, {
-                         'children': [{'children': [{'id': 14, 'label': '<Node (14)>'}], 'id': 13, 'label': '<Node (13)>'}, {'children': [{'id': 16, 'label': '<Node (16)>'}, {'id': 17, 'label': '<Node (17)>'}], 'id': 15, 'label': '<Node (15)>'}, {'children': [{'children': [{'id': 20, 'label': '<Node (20)>'}], 'id': 19, 'label': '<Node (19)>'}, {'children': [{'id': 22, 'label': '<Node (22)>'}], 'id': 21, 'label': '<Node (21)>'}], 'id': 18, 'label': '<Node (18)>'}], 'id': 12, 'label': '<Node (12)>'}])
+        tree_reqursively = self.model.get_tree_reqursively(self.session,
+                                                           json=True)
+        self.assertEqual(tree, reference_tree)
+        self.assertEqual(tree_reqursively, reference_tree)
 
     def test_get_json_tree_with_custom_field(self):
         """.. note::
@@ -1292,9 +1354,14 @@ def fields(node):
         """
         def fields(node):
             return {'visible': node.visible}
+
+        reference_tree = [{'visible': None, 'children': [{'visible': True, 'children': [{'visible': True, 'id': 3, 'label': '<Node (3)>'}], 'id': 2, 'label': '<Node (2)>'}, {'visible': True, 'children': [{'visible': True, 'id': 5, 'label': '<Node (5)>'}, {'visible': True, 'id': 6, 'label': '<Node (6)>'}], 'id': 4, 'label': '<Node (4)>'}, {'visible': True, 'children': [{'visible': True, 'children': [{'visible': None, 'id': 9, 'label': '<Node (9)>'}], 'id': 8, 'label': '<Node (8)>'}, {'visible': None, 'children': [{'visible': None, 'id': 11, 'label': '<Node (11)>'}], 'id': 10, 'label': '<Node (10)>'}], 'id': 7, 'label': '<Node (7)>'}], 'id': 1, 'label': '<Node (1)>'}, {'visible': None, 'children': [{'visible': None, 'children': [{'visible': None, 'id': 14, 'label': '<Node (14)>'}], 'id': 13, 'label': '<Node (13)>'}, {'visible': None, 'children': [{'visible': None, 'id': 16, 'label': '<Node (16)>'}, {'visible': None, 'id': 17, 'label': '<Node (17)>'}], 'id': 15, 'label': '<Node (15)>'}, {'visible': None, 'children': [{'visible': None, 'children': [{'visible': None, 'id': 20, 'label': '<Node (20)>'}], 'id': 19, 'label': '<Node (19)>'}, {'visible': None, 'children': [{'visible': None, 'id': 22, 'label': '<Node (22)>'}], 'id': 21, 'label': '<Node (21)>'}], 'id': 18, 'label': '<Node (18)>'}], 'id': 12, 'label': '<Node (12)>'}]  # noqa
+
         tree = self.model.get_tree(self.session, json=True, json_fields=fields)
-        self.assertEqual(tree, [{'visible': None, 'children': [{'visible': True, 'children': [{'visible': True, 'id': 3, 'label': '<Node (3)>'}], 'id': 2, 'label': '<Node (2)>'}, {'visible': True, 'children': [{'visible': True, 'id': 5, 'label': '<Node (5)>'}, {'visible': True, 'id': 6, 'label': '<Node (6)>'}], 'id': 4, 'label': '<Node (4)>'}, {'visible': True, 'children': [{'visible': True, 'children': [{'visible': None, 'id': 9, 'label': '<Node (9)>'}], 'id': 8, 'label': '<Node (8)>'}, {'visible': None, 'children': [{'visible': None, 'id': 11, 'label': '<Node (11)>'}], 'id': 10, 'label': '<Node (10)>'}], 'id': 7, 'label': '<Node (7)>'}], 'id': 1, 'label': '<Node (1)>'}, {
-                         'visible': None, 'children': [{'visible': None, 'children': [{'visible': None, 'id': 14, 'label': '<Node (14)>'}], 'id': 13, 'label': '<Node (13)>'}, {'visible': None, 'children': [{'visible': None, 'id': 16, 'label': '<Node (16)>'}, {'visible': None, 'id': 17, 'label': '<Node (17)>'}], 'id': 15, 'label': '<Node (15)>'}, {'visible': None, 'children': [{'visible': None, 'children': [{'visible': None, 'id': 20, 'label': '<Node (20)>'}], 'id': 19, 'label': '<Node (19)>'}, {'visible': None, 'children': [{'visible': None, 'id': 22, 'label': '<Node (22)>'}], 'id': 21, 'label': '<Node (21)>'}], 'id': 18, 'label': '<Node (18)>'}], 'id': 12, 'label': '<Node (12)>'}])
+        tree_reqursively = self.model.get_tree(self.session, json=True,
+                                               json_fields=fields)
+        self.assertEqual(tree, reference_tree)
+        self.assertEqual(tree_reqursively, reference_tree)
 
     def test_rebuild(self):
         """ Rebuild tree with tree_id==1