Refactor With nodes to have subnodes, WithItems, corresponding to the stdlib 3 With node.

6 files changed, 57 insertions, 40 deletions

diff --git a/astroid/as_string.py b/astroid/as_string.py
index eefc3c47..7f59f88c 100644
--- a/astroid/as_string.py
+++ b/astroid/as_string.py
@@ -406,11 +406,13 @@ class AsStringVisitor(object):
 
     def visit_with(self, node): # 'with' without 'as' is possible
         """return an astroid.With node as string"""
-        items = ', '.join(('(%s)' % expr.accept(self)) +
-                          (vars and ' as (%s)' % (vars.accept(self)) or '')
-                          for expr, vars in node.items)
+        items = ', '.join(item.accept(self) for item in node.items)
         return 'with %s:\n%s' % (items, self._stmt_list(node.body))
 
+    def visit_withitem(self, node):
+        return ('(%s)' % node.context_expr.accept(self) + ' as (%s)' %
+                (optional_vars.accept(self)) if optional_vars else '')
+
     def visit_yield(self, node):
         """yield an ast.Yield node as string"""
         yi_val = node.value and (" " + node.value.accept(self)) or ""
diff --git a/astroid/nodes.py b/astroid/nodes.py
index e2fb1e0c..1aea53f2 100644
--- a/astroid/nodes.py
+++ b/astroid/nodes.py
@@ -44,7 +44,7 @@ from astroid.tree.node_classes import (
     ImportFrom, Attribute, Global, If, IfExp, Import, Index, Keyword,
     List, Name, NameConstant, Nonlocal, Pass, Print, Raise, Return, Set, Slice,
     Starred, Subscript, TryExcept, TryFinally, Tuple, UnaryOp, While, With,
-    Yield, YieldFrom, AsyncFor, Await, AsyncWith,
+    WithItem, Yield, YieldFrom, AsyncFor, Await, AsyncWith,
     # Node not present in the builtin ast module.
     DictUnpack,
     # Special nodes for building from live objects.
@@ -79,6 +79,6 @@ ALL_NODE_CLASSES = (
     Set, SetComp, Slice, Starred, Subscript,
     TryExcept, TryFinally, Tuple,
     UnaryOp, Unknown,
-    While, With,
+    While, With, WithItem,
     Yield, YieldFrom,
     )
diff --git a/astroid/protocols.py b/astroid/protocols.py
index c9e2842d..0464f7e7 100644
--- a/astroid/protocols.py
+++ b/astroid/protocols.py
@@ -469,6 +469,12 @@ def _infer_context_manager(self, mgr, context, nodes):
             yield result
 
 
+@assigned_stmts.register(treeabc.WithItem)
+def withitem_assigned_stmts(self, nodes, node=None, context=None, asspath=None):
+    print(node, self, self.parent)
+    return self.parent.assigned_stmts(nodes, node=node, asspath=asspath)
+
+
 @assigned_stmts.register(treeabc.With)
 @decorators.raise_if_nothing_inferred
 def with_assigned_stmts(self, nodes, node=None, context=None, asspath=None):
@@ -495,7 +501,11 @@ def with_assigned_stmts(self, nodes, node=None, context=None, asspath=None):
         context: TODO
         asspath: TODO
     """
-    mgr = next(mgr for (mgr, vars) in self.items if vars == node)
+    print('SELF', self)
+    print('NODES', nodes)
+    print('NODE', node)
+    mgr = next(item.context_expr for item in self.items if
+               item.optional_vars == node) 
     if asspath is None:
         for result in _infer_context_manager(self, mgr, context, nodes):
             yield result
diff --git a/astroid/tree/node_classes.py b/astroid/tree/node_classes.py
index 06386226..ed757632 100644
--- a/astroid/tree/node_classes.py
+++ b/astroid/tree/node_classes.py
@@ -1268,8 +1268,11 @@ class With(base.BlockRangeMixIn, base.AssignTypeMixin,
            AssignedStmtsMixin, Statement):
     """class representing a With node"""
     _astroid_fields = ('items', 'body')
-    items = None
-    body = None
+
+    def __init__(self, lineno=None, col_offset=None, parent=None):
+        self.items = []
+        self.body = []
+        super(With, self).__init__(lineno, col_offset, parent)
 
     def postinit(self, items=None, body=None):
         self.items = items
@@ -1277,15 +1280,18 @@ class With(base.BlockRangeMixIn, base.AssignTypeMixin,
 
     @decorators.cachedproperty
     def blockstart_tolineno(self):
-        return self.items[-1][0].tolineno
+        return self.items[-1].context_expr.tolineno
 
-    def get_children(self):
-        for expr, var in self.items:
-            yield expr
-            if var:
-                yield var
-        for elt in self.body:
-            yield elt
+
+@util.register_implementation(treeabc.WithItem)
+class WithItem(base.NodeNG, base.ParentAssignTypeMixin, AssignedStmtsMixin):
+    _astroid_fields = ('context_expr', 'optional_vars')
+    context_expr = None
+    optional_vars = None
+
+    def postinit(self, context_expr=None, optional_vars=None):
+        self.context_expr = context_expr
+        self.optional_vars = optional_vars
 
 
 @util.register_implementation(treeabc.AsyncWith)
diff --git a/astroid/tree/rebuilder.py b/astroid/tree/rebuilder.py
index 0221b190..12b9388c 100644
--- a/astroid/tree/rebuilder.py
+++ b/astroid/tree/rebuilder.py
@@ -709,14 +709,13 @@ class TreeRebuilder(object):
 
     def visit_with(self, node, parent):
         newnode = nodes.With(node.lineno, node.col_offset, parent)
-        expr = self.visit(node.context_expr, newnode)
-        if node.optional_vars is not None:
-            optional_vars = self.visit(node.optional_vars, newnode)
-        else:
-            optional_vars = None
-        newnode.postinit([(expr, optional_vars)],
-                         [self.visit(child, newnode)
-                          for child in node.body])
+        with_item = nodes.WithItem(node.context_expr.lineno,
+                                   node.context_expr.col_offset, newnode)
+        context_expr = self.visit(node.context_expr, with_item)
+        optional_vars = _visit_or_none(node, 'optional_vars', self, with_item)
+        with_item.postinit(context_expr, optional_vars)
+        newnode.postinit([with_item],
+                         [self.visit(child, newnode) for child in node.body])
         return newnode
 
     def visit_yield(self, node, parent):
@@ -793,23 +792,19 @@ class TreeRebuilder3(TreeRebuilder):
         elif node.handlers:
             return self.visit_tryexcept(node, parent)
 
-    def _visit_with(self, cls, node, parent):
-        if 'items' not in node._fields:
-            # python < 3.3
-            return super(TreeRebuilder3, self).visit_with(node, parent)
-
-        newnode = cls(node.lineno, node.col_offset, parent)
-        def visit_child(child):
-            expr = self.visit(child.context_expr, newnode)
-            var = _visit_or_none(child, 'optional_vars', self, newnode)
-            return expr, var
-        newnode.postinit([visit_child(child) for child in node.items],
-                         [self.visit(child, newnode)
-                          for child in node.body])
+    def visit_with(self, node, parent):
+        newnode = nodes.With(node.lineno, node.col_offset, parent)
+        newnode.postinit([self.visit(item, newnode) for item in node.items],
+                         [self.visit(child, newnode) for child in node.body])
         return newnode
 
-    def visit_with(self, node, parent):
-        return self._visit_with(nodes.With, node, parent)
+    def visit_withitem(self, node, parent):
+        newnode = nodes.WithItem(node.context_expr.lineno,
+                                 node.context_expr.col_offset, parent)
+        context_expr = self.visit(node.context_expr, newnode)
+        optional_vars = _visit_or_none(node, 'optional_vars', self, newnode)
+        newnode.postinit(context_expr=context_expr, optional_vars=optional_vars)
+        return newnode
 
     def visit_yieldfrom(self, node, parent):
         newnode = nodes.YieldFrom(node.lineno, node.col_offset, parent)
@@ -834,7 +829,7 @@ class TreeRebuilder3(TreeRebuilder):
         return newnode
 
     def visit_asyncwith(self, node, parent):
-        return self._visit_with(nodes.AsyncWith, node, parent)
+        return self.visit_with(nodes.AsyncWith, node, parent)
 
 
 if sys.version_info >= (3, 0):
diff --git a/astroid/tree/treeabc.py b/astroid/tree/treeabc.py
index c9fa2269..6ad96ba8 100644
--- a/astroid/tree/treeabc.py
+++ b/astroid/tree/treeabc.py
@@ -273,6 +273,10 @@ class With(Statement):
     """Class representing a With node"""
 
 
+class WithItem(NodeNG):
+    """Class representing a WithItem node"""
+
+
 class AsyncWith(With):
     """Asynchronous `with` built with the `async` keyword."""