Commit 4f0f0a00 authored by Xavier Thompson's avatar Xavier Thompson

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

parent e9d9d2d8
...@@ -178,6 +178,15 @@ class TestTreeFragments(CythonTest): ...@@ -178,6 +178,15 @@ class TestTreeFragments(CythonTest):
self.assertTrue(isinstance(T.stats[0], PropertyNode), T) self.assertTrue(isinstance(T.stats[0], PropertyNode), T)
self.assertTrue(T.stats[0].name == "test", 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__": if __name__ == "__main__":
import unittest import unittest
unittest.main() unittest.main()
...@@ -246,6 +246,13 @@ class TemplateTransform(VisitorTransform): ...@@ -246,6 +246,13 @@ class TemplateTransform(VisitorTransform):
else: else:
return self.visit_Node(node) 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): def copy_code_tree(node):
return TreeCopier()(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