Commit 45e1d7e0 authored by Xavier Thompson's avatar Xavier Thompson

Allow TreeFragments to substitute 'as' target names in 'from ... import ... as ...' statements

parent bca614e2
......@@ -178,6 +178,15 @@ class TestTreeFragments(CythonTest):
self.assertTrue(isinstance(T.stats[0], PropertyNode), T)
self.assertTrue(T.stats[0].name == "test", T)
def test_fromimportstatnode(self):
F = TreeFragment(u"from A import b as NAME")
T = F.substitute({
"NAME" : NameNode(None, name="test"),
})
self.assertTrue(isinstance(T.stats[0], FromImportStatNode), T)
self.assertTrue(isinstance(T.stats[0].items[0][1], NameNode), T)
self.assertTrue(T.stats[0].items[0][1].name == "test", T)
if __name__ == "__main__":
import unittest
unittest.main()
......@@ -246,6 +246,13 @@ class TemplateTransform(VisitorTransform):
else:
return self.visit_Node(node)
def visit_FromImportStatNode(self, node):
for i, item in enumerate(node.items):
name, target_namenode = item
new_namenode = self.visit(target_namenode)
node.items[i] = (name, new_namenode)
return self.visit_Node(node)
def copy_code_tree(node):
return TreeCopier()(node)
......
Markdown is supported
0%
or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment