Improve zipper's backwards compatibility and fix bugs

4 files changed, 37 insertions, 11 deletions

diff --git a/astroid/__init__.py b/astroid/__init__.py
index e32e6ba0..7d84335e 100644
--- a/astroid/__init__.py
+++ b/astroid/__init__.py
@@ -82,8 +82,7 @@ from astroid.interpreter.lookup import builtin_lookup
 from astroid.builder import parse
 from astroid.util import Uninferable, YES
 
-# ???
-from astroid.tree import zipper
+# TODO
 from astroid.tree.base import NodeNG
 
 # transform utilities (filters and decorator)
diff --git a/astroid/test_utils.py b/astroid/test_utils.py
index ff81cd59..c1ff5ebd 100644
--- a/astroid/test_utils.py
+++ b/astroid/test_utils.py
@@ -37,6 +37,7 @@ def _extract_expressions(node):
             and isinstance(node.func, nodes.Name)
             and node.func.name == _TRANSIENT_FUNCTION):
         real_expr = node.args[0]
+        # real_expr = node.down().right().down()
         real_expr.parent = node.parent
         # Search for node in all _astng_fields (the fields checked when
         # get_children is called) of its parent. Some of those fields may
diff --git a/astroid/tests/unittest_zipper.py b/astroid/tests/unittest_zipper.py
index 1c2cca17..a10b7179 100644
--- a/astroid/tests/unittest_zipper.py
+++ b/astroid/tests/unittest_zipper.py
@@ -83,7 +83,7 @@ def ast_from_file_name(name):
         return AST_CACHE[name]
     with open(name, 'r') as source_file:
         # print(name)
-        root = astroid.parse(source_file.read())
+        root = astroid.parse(source_file.read()).__wrapped__
         ast = ASTMap()
         AssignLabels()(ast, root)
         to_visit = [1]
@@ -137,7 +137,7 @@ def postorder_descendants(label, ast, dont_recurse_on=None):
 
 # This test function uses a set-based implementation for finding the
 # common parent rather than the reverse-based implementation in the
-# functional code.
+# actual code.
 def common_ancestor(label1, label2, ast):
     ancestors = set()
     while label1:
@@ -162,6 +162,14 @@ def traverse_to_node(label, ast, location):
         location = move(location)
     return location
 
+def get_children(label, ast):
+    for child_label in ast[label].children:
+        if isinstance(ast[child_label].node, collections.Sequence):
+            for grandchild_label in ast[child_label].children:
+                yield grandchild_label
+        else:
+            yield child_label
+
 # This function and strategy creates a strategy for generating a
 # random node class, for testing that the iterators properly exclude
 # nodes of that type and their descendants.
@@ -213,7 +221,7 @@ class TestZipper(unittest.TestCase):
         nodes = tuple(ast)
         random_label = choice(nodes)
         random_node = zipper.Zipper(ast[random_label].node)
-        for node, label in zip(random_node.get_children(), ast[random_label].children):
+        for node, label in zip(random_node.children(), ast[random_label].children):
             self.assertIs(node.__wrapped__, ast[label].node)
         for node, label in zip(random_node.preorder_descendants(), preorder_descendants(random_label, ast)):
             self.assertIs(node.__wrapped__, ast[label].node)
@@ -223,6 +231,8 @@ class TestZipper(unittest.TestCase):
             self.assertIs(node.__wrapped__, ast[label].node)
         for node, label in zip(random_node.postorder_descendants(dont_recurse_on=node_type), postorder_descendants(random_label, ast, dont_recurse_on=node_type)):
             self.assertIs(node.__wrapped__, ast[label].node)
+        for node, label in zip(random_node.get_children(), get_children(random_label, ast)):
+            self.assertIs(node.__wrapped__, ast[label].node)
 
     @hypothesis.settings(perform_health_check=False)
     @hypothesis.given(ast_strategy, strategies.choices())
diff --git a/astroid/tree/zipper.py b/astroid/tree/zipper.py
index cd30b6e9..51360851 100644
--- a/astroid/tree/zipper.py
+++ b/astroid/tree/zipper.py
@@ -260,7 +260,7 @@ class Zipper(wrapt.ObjectProxy):
                 location = location.up()
         return location
 
-    def get_children(self):
+    def children(self):
         '''Iterates over the children of the focus.'''
         child = self.down()
         while child is not None:
@@ -285,10 +285,10 @@ class Zipper(wrapt.ObjectProxy):
             yield location
             if dont_recurse_on is None:
                 to_visit.extend(c for c in
-                                reversed(tuple(location.get_children())))
+                                reversed(tuple(location.children())))
             else:
                 to_visit.extend(c for c in
-                                reversed(tuple(location.get_children()))
+                                reversed(tuple(location.children()))
                                 if not isinstance(c, dont_recurse_on))
 
     def postorder_descendants(self, dont_recurse_on=None):
@@ -306,10 +306,10 @@ class Zipper(wrapt.ObjectProxy):
                 visited_ancestors.append(location)
                 if dont_recurse_on is None:
                     to_visit.extend(c for c in
-                                    reversed(tuple(location.get_children())))
+                                    reversed(tuple(location.children())))
                 else:
                     to_visit.extend(c for c in
-                                    reversed(tuple(location.get_children()))
+                                    reversed(tuple(location.children()))
                                     if not isinstance(c, dont_recurse_on))
                 continue
             visited_ancestors.pop()
@@ -325,7 +325,7 @@ class Zipper(wrapt.ObjectProxy):
                 not include nodes of this type or types or any of the
                 descendants of those nodes.
         '''
-        return (d for d in self.preorder_descendants(skip_class) if isinstance(node, cls))
+        return (d for d in self.preorder_descendants(skip_class) if isinstance(d, cls))
         # if isinstance(self, cls):
         #     yield self
         # child = self.down()
@@ -350,6 +350,22 @@ class Zipper(wrapt.ObjectProxy):
         else:
             return location
 
+    def get_children(self):
+        '''Iterates over nodes that are children or grandchildren, no
+        sequences.
+
+        '''
+        child = self.down()
+        while child is not None:
+            if isinstance(child, collections.Sequence):
+                grandchild = child.down()
+                for _ in range(len(child)):
+                    yield grandchild
+                    grandchild = grandchild.right()
+            else:
+                yield child
+            child = child.right()
+
     def last_child(self):
         return self.rightmost()