diff options
| -rw-r--r-- | _nodes_ast.py | 11 | ||||
| -rw-r--r-- | node_classes.py | 8 | ||||
| -rw-r--r-- | nodes.py | 6 | ||||
| -rw-r--r-- | nodes_as_string.py | 9 | ||||
| -rw-r--r-- | test/unittest_lookup.py | 20 | ||||
| -rw-r--r-- | utils.py | 3 |
6 files changed, 52 insertions, 5 deletions
diff --git a/_nodes_ast.py b/_nodes_ast.py index 4b6767f0..4dbca6d0 100644 --- a/_nodes_ast.py +++ b/_nodes_ast.py @@ -328,6 +328,17 @@ class TreeRebuilder(RebuildVisitor): newnode.set_line_info(newnode.last_child()) return newnode + def visit_dictcomp(self, node, parent): + """visit a DictComp node by returning a fresh instance of it""" + newnode = new.DictComp() + _lineno_parent(node, newnode, parent) + newnode.key = self.visit(node.key, newnode) + newnode.value = self.visit(node.value, newnode) + newnode.generators = [self.visit(child, newnode) + for child in node.generators] + newnode.set_line_info(newnode.last_child()) + return newnode + def visit_discard(self, node, parent): """visit a Discard node by returning a fresh instance of it""" newnode = new.Discard() diff --git a/node_classes.py b/node_classes.py index c9311759..e93330ba 100644 --- a/node_classes.py +++ b/node_classes.py @@ -505,6 +505,14 @@ class Dict(NodeNG, Instance): raise IndexError(key) +class DictComp(NodeNG): + """class representing a DictComp node""" + _astng_fields = ('key', 'value', 'generators') + key = None + value = None + generators = None + + class Discard(StmtMixIn, NodeNG): """class representing a Discard node""" _astng_fields = ('value',) @@ -54,8 +54,8 @@ __docformat__ = "restructuredtext en" from logilab.astng.node_classes import Arguments, AssAttr, Assert, Assign, \ AssName, AugAssign, Backquote, BinOp, BoolOp, Break, CallFunc, Compare, \ Comprehension, Const, Continue, Decorators, DelAttr, DelName, Delete, \ - Dict, Discard, Ellipsis, EmptyNode, ExceptHandler, Exec, ExtSlice, For, \ - From, Getattr, Global, If, IfExp, Import, Index, Keyword, \ + Dict, DictComp, Discard, Ellipsis, EmptyNode, ExceptHandler, Exec, \ + ExtSlice, For, From, Getattr, Global, If, IfExp, Import, Index, Keyword, \ List, ListComp, Name, Pass, Print, Raise, Return, Slice, Subscript, \ TryExcept, TryFinally, Tuple, UnaryOp, While, With, Yield, const_factory from logilab.astng.scoped_nodes import Module, GenExpr, Lambda, Function, Class @@ -65,7 +65,7 @@ ALL_NODE_CLASSES = ( Backquote, BinOp, BoolOp, Break, CallFunc, Class, Compare, Comprehension, Const, Continue, Decorators, DelAttr, DelName, Delete, - Dict, Discard, + Dict, DictComp, Discard, Ellipsis, EmptyNode, ExceptHandler, Exec, ExtSlice, For, From, Function, Getattr, GenExpr, Global, diff --git a/nodes_as_string.py b/nodes_as_string.py index 0a0f287f..dc4db2da 100644 --- a/nodes_as_string.py +++ b/nodes_as_string.py @@ -173,7 +173,12 @@ class AsStringVisitor(ASTVisitor): """return an astng.Dict node as string""" return '{%s}' % ', '.join(['%s: %s' % (key.accept(self), value.accept(self)) for key, value in node.items]) - + + def visit_dictcomp(self, node): + """return an astng.DictComp node as string""" + return '[%s: %s %s]' % (node.key.accept(self), node.value.accept(self), + ' '.join([n.accept(self) for n in node.generators])) + def visit_discard(self, node): """return an astng.Discard node as string""" return node.value.accept(self) @@ -233,7 +238,7 @@ class AsStringVisitor(ASTVisitor): docs, self._stmt_list(node.body)) def visit_genexpr(self, node): - """return an astng.ListComp node as string""" + """return an astng.GenExpr node as string""" return '(%s %s)' % (node.elt.accept(self), ' '.join([n.accept(self) for n in node.generators])) diff --git a/test/unittest_lookup.py b/test/unittest_lookup.py index ae6d9bc9..dc0966f0 100644 --- a/test/unittest_lookup.py +++ b/test/unittest_lookup.py @@ -169,6 +169,26 @@ print list( i for i in range(10) ) self.assertEqual(xnames[2].lookup('i')[1][0].lineno, 4) + def test_dict_comps(self): + if sys.version_info < (2, 7): + self.skipTest('this test require python >= 2.7') + astng = builder.string_build(""" +print { i: j for i in range(10) for j in range(10) } +print { i: j for i in range(10) for j in range(10) } + """, __name__, __file__) + xnames = [n for n in astng.nodes_of_class(nodes.Name) if n.name == 'i'] + self.assertEqual(len(xnames[0].lookup('i')[1]), 1) + self.assertEqual(xnames[0].lookup('i')[1][0].lineno, 2) + self.assertEqual(len(xnames[1].lookup('i')[1]), 1) + self.assertEqual(xnames[1].lookup('i')[1][0].lineno, 3) + + xnames = [n for n in astng.nodes_of_class(nodes.Name) if n.name == 'j'] + self.assertEqual(len(xnames[0].lookup('i')[1]), 1) + self.assertEqual(xnames[0].lookup('i')[1][0].lineno, 2) + self.assertEqual(len(xnames[1].lookup('i')[1]), 1) + self.assertEqual(xnames[1].lookup('i')[1][0].lineno, 3) + + def test_explicit___name__(self): code = ''' class Pouet: @@ -114,6 +114,9 @@ class ASTVisitor(object): def visit_dict(self, node): """dummy method for visiting an Dict node""" + def visit_dictcomp(self, node): + """dummy method for visiting an DictComp node""" + def visit_discard(self, node): """dummy method for visiting an Discard node""" |