ParentNodeTransformer is now ParentChildNodeTransformer and adds link…

…s to node children

README.rst

@@ -23,7 +23,7 @@ If you want to have latest changes you should clone this repository and use ``se
     $ cd astmonkey
     $ python setup.py install
 
-transformers.ParentNodeTransformer
+transformers.ParentChildNodeTransformer
 ----------------------------------
 
 This transformer adds few fields to every node in AST:
@@ -32,6 +32,7 @@ This transformer adds few fields to every node in AST:
 * ``parents`` - list of all parents (only ``ast.expr_context`` nodes have more than one parent node, in other causes this is one-element list),
 * ``parent_field`` - name of field in parent node including child node,
 * ``parent_field_index`` - parent node field index, if it is a list.
+* ``children`` - link to children nodes.
 
 Example usage:
 
@@ -41,18 +42,19 @@ Example usage:
     from astmonkey import transformers
 
     node = ast.parse('x = 1')
-    node = transformers.ParentNodeTransformer().visit(node)
+    node = transformers.ParentChildNodeTransformer().visit(node)
 
     assert(node == node.body[0].parent)
     assert(node.body[0].parent_field == 'body')
     assert(node.body[0].parent_field_index == 0)
+    assert(node.body[0] in node.children)
 
 visitors.GraphNodeVisitor
 -------------------------
 
 This visitor creates Graphviz graph from Python AST (via ``pydot``). Before you use 
 ``GraphNodeVisitor`` you need to add parents links to tree nodes (with 
-``ParentNodeTransformer``).
+``ParentChildNodeTransformer``).
 
 Example usage:
 
@@ -62,7 +64,7 @@ Example usage:
     from astmonkey import visitors, transformers
 
     node = ast.parse('def foo(x):\n\treturn x + 1')
-    node = transformers.ParentNodeTransformer().visit(node)
+    node = transformers.ParentChildNodeTransformer().visit(node)
     visitor = visitors.GraphNodeVisitor()
     visitor.visit(node)
 
@@ -98,7 +100,7 @@ utils.is_docstring
 
 This routine checks if target node is a docstring. Before you use 
 ``is_docstring`` you need to add parents links to tree nodes (with 
-``ParentNodeTransformer``).
+``ParentChildNodeTransformer``).
 
 Example usage:
 
@@ -108,7 +110,7 @@ Example usage:
     from astmonkey import utils, transformers
 
     node = ast.parse('def foo(x):\n\t"""doc"""')
-    node = transformers.ParentNodeTransformer().visit(node)
+    node = transformers.ParentChildNodeTransformer().visit(node)
 
     docstring_node = node.body[0].body[0].value
     assert(not utils.is_docstring(node))

astmonkey/tests/test_transformers.py

@@ -9,19 +9,20 @@
 from astmonkey import transformers
 
 
-class TestParentNodeTransformer(object):
+class TestParentChildNodeTransformer(object):
     @pytest.fixture
     def transformer(self):
-        return transformers.ParentNodeTransformer()
+        return transformers.ParentChildNodeTransformer()
 
-    def test_module_node_parent(self, transformer):
+    def test_module_node(self, transformer):
         node = ast.parse('')
 
         transformed_node = transformer.visit(node)
 
         assert transformed_node.parent is None
+        assert transformed_node.children == []
 
-    def test_non_module_node_parent(self, transformer):
+    def test_non_module_node(self, transformer):
         node = ast.parse('x = 1')
 
         transformed_node = transformer.visit(node)
@@ -30,8 +31,10 @@ def test_non_module_node_parent(self, transformer):
         assert transformed_node == assign_node.parent
         assert assign_node.parent_field == 'body'
         assert assign_node.parent_field_index == 0
+        assert transformed_node.children == [assign_node]
+        assert len(assign_node.children) == 2
 
-    def test_expr_context_nodes_parent(self, transformer):
+    def test_expr_context_nodes(self, transformer):
         node = ast.parse('x = 1\nx = 2')
 
         transformer.visit(node)
@@ -41,3 +44,5 @@ def test_expr_context_nodes_parent(self, transformer):
         second_name_node = node.body[1].targets[0]
         assert first_name_node in ctx_node.parents
         assert second_name_node in ctx_node.parents
+        assert ctx_node in first_name_node.children
+        assert ctx_node in second_name_node.children

astmonkey/tests/test_utils.py

@@ -6,21 +6,21 @@
 
 class TestIsDocstring(unittest.TestCase):
     def test_non_docstring_node(self):
-        node = transformers.ParentNodeTransformer().visit(ast.parse(''))
+        node = transformers.ParentChildNodeTransformer().visit(ast.parse(''))
 
         assert not utils.is_docstring(node)
 
     def test_module_docstring_node(self):
-        node = transformers.ParentNodeTransformer().visit(ast.parse('"""doc"""'))
+        node = transformers.ParentChildNodeTransformer().visit(ast.parse('"""doc"""'))
 
         assert utils.is_docstring(node.body[0].value)
 
     def test_function_docstring_node(self):
-        node = transformers.ParentNodeTransformer().visit(ast.parse('def foo():\n\t"""doc"""'))
+        node = transformers.ParentChildNodeTransformer().visit(ast.parse('def foo():\n\t"""doc"""'))
 
         assert utils.is_docstring(node.body[0].body[0].value)
 
     def test_class_docstring_node(self):
-        node = transformers.ParentNodeTransformer().visit(ast.parse('class X:\n\t"""doc"""'))
+        node = transformers.ParentChildNodeTransformer().visit(ast.parse('class X:\n\t"""doc"""'))
 
         assert utils.is_docstring(node.body[0].body[0].value)

astmonkey/tests/test_visitors.py

@@ -17,37 +17,37 @@ def visitor(self):
         return visitors.GraphNodeVisitor()
 
     def test_has_edge(self, visitor):
-        node = transformers.ParentNodeTransformer().visit(ast.parse('x = 1'))
+        node = transformers.ParentChildNodeTransformer().visit(ast.parse('x = 1'))
 
         visitor.visit(node)
 
         assert visitor.graph.get_edge(str(node), str(node.body[0]))
 
     def test_does_not_have_edge(self, visitor):
-        node = transformers.ParentNodeTransformer().visit(ast.parse('x = 1'))
+        node = transformers.ParentChildNodeTransformer().visit(ast.parse('x = 1'))
 
         visitor.visit(node)
 
         assert not visitor.graph.get_edge(str(node), str(node.body[0].value))
 
     def test_node_label(self, visitor):
-        node = transformers.ParentNodeTransformer().visit(ast.parse('x = 1'))
+        node = transformers.ParentChildNodeTransformer().visit(ast.parse('x = 1'))
 
         visitor.visit(node)
 
         dot_node = visitor.graph.get_node(str(node.body[0].value))[0]
         assert dot_node.get_label() == 'ast.Num(n=1)'
 
     def test_edge_label(self, visitor):
-        node = transformers.ParentNodeTransformer().visit(ast.parse('x = 1'))
+        node = transformers.ParentChildNodeTransformer().visit(ast.parse('x = 1'))
 
         visitor.visit(node)
 
         dot_edge = visitor.graph.get_edge(str(node), str(node.body[0]))[0]
         assert dot_edge.get_label() == 'body[0]'
 
     def test_multi_parents_node_label(self, visitor):
-        node = transformers.ParentNodeTransformer().visit(ast.parse('x = 1\nx = 2'))
+        node = transformers.ParentChildNodeTransformer().visit(ast.parse('x = 1\nx = 2'))
 
         visitor.visit(node)
 

astmonkey/transformers.py

@@ -1,26 +1,34 @@
 import ast
 
 
-class ParentNodeTransformer(object):
+class ParentChildNodeTransformer(object):
 
     def visit(self, node):
+        self._prepare_node(node)
+        for field, value in ast.iter_fields(node):
+            self._process_field(node, field, value)
+        return node
+
+    @staticmethod
+    def _prepare_node(node):
         if not hasattr(node, 'parent'):
             node.parent = None
             node.parents = []
-        for field, value in ast.iter_fields(node):
-            if isinstance(value, list):
-                for index, item in enumerate(value):
-                    if isinstance(item, ast.AST):
-                        self.visit(item)
-                        self._set_parnt_fields(item, node, field, index)
-            elif isinstance(value, ast.AST):
-                self.visit(value)
-                self._set_parnt_fields(value, node, field)
-        return node
+        if not hasattr(node, 'children'):
+            node.children = []
 
-    def _set_parnt_fields(self, node, parent, field, index=None):
-        node.parent = parent
-        node.parents.append(parent)
-        node.parent_field = field
-        node.parent_field_index = index
+    def _process_field(self, node, field, value):
+        if isinstance(value, list):
+            for index, item in enumerate(value):
+                if isinstance(item, ast.AST):
+                    self._process_child(item, node, field, index)
+        elif isinstance(value, ast.AST):
+            self._process_child(value, node, field)
 
+    def _process_child(self, child, parent, field_name, index=None):
+        self.visit(child)
+        child.parent = parent
+        child.parents.append(parent)
+        child.parent_field = field_name
+        child.parent_field_index = index
+        child.parent.children.append(child)

examples/graph_node_visitor.py

@@ -2,7 +2,7 @@
 from astmonkey import visitors, transformers
 
 node = ast.parse('def foo(x):\n\treturn x + 1')
-node = transformers.ParentNodeTransformer().visit(node)
+node = transformers.ParentChildNodeTransformer().visit(node)
 visitor = visitors.GraphNodeVisitor()
 visitor.visit(node)
 

examples/is_docstring.py

@@ -2,7 +2,7 @@
 from astmonkey import utils, transformers
 
 node = ast.parse('def foo(x):\n\t"""doc"""')
-node = transformers.ParentNodeTransformer().visit(node)
+node = transformers.ParentChildNodeTransformer().visit(node)
 
 docstring_node = node.body[0].body[0].value
 assert(not utils.is_docstring(node))

examples/parent_node_transformer.py

@@ -2,7 +2,7 @@
 from astmonkey import transformers
 
 node = ast.parse('x = 1')
-node = transformers.ParentNodeTransformer().visit(node)
+node = transformers.ParentChildNodeTransformer().visit(node)
 
 assert(node == node.body[0].parent)
 assert(node.body[0].parent_field == 'body')