TreeFragment fix: Replace enclosing ExprStatNode if statement is substituted

diff --git a/Cython/Compiler/Tests/TestTreeFragment.py b/Cython/Compiler/Tests/TestTreeFragment.py
index 7e283b9d695d8c118e7b55728b06f130420c7350..1658398f5d7253724ff919bb37eb42ac214c0005 100644 (file)
--- a/Cython/Compiler/Tests/TestTreeFragment.py
+++ b/Cython/Compiler/Tests/TestTreeFragment.py
@@ -1,5 +1,6 @@
 from Cython.TestUtils import CythonTest
 from Cython.Compiler.TreeFragment import *
+from Cython.Compiler.Nodes import *
 
 class TestTreeFragments(CythonTest):
     def test_basic(self):
@@ -21,6 +22,12 @@ class TestTreeFragments(CythonTest):
         T = F.substitute({"x" : y})
         self.assertCode(u"y = 4", T)
 
+    def test_exprstat(self):
+        F = self.fragment(u"PASS")
+        pass_stat = PassStatNode(pos=None)
+        T = F.substitute({"PASS" : pass_stat})
+        self.assert_(T.body is pass_stat, T.body)
+
 if __name__ == "__main__":
     import unittest
     unittest.main()

diff --git a/Cython/Compiler/TreeFragment.py b/Cython/Compiler/TreeFragment.py
index 8feab2f00a76bb1a7833c5490dd7242bda2c9570..9db389f1d2e8bbd30b8980ab1b12c2f738cc0e78 100644 (file)
--- a/Cython/Compiler/TreeFragment.py
+++ b/Cython/Compiler/TreeFragment.py
@@ -83,6 +83,14 @@ class SubstitutionTransform(VisitorTransform):
             # Clone
             return self.visit_Node(node)
     
+    def visit_ExprStatNode(self, node):
+        # If an expression-as-statement consists of only a replaceable
+        # NameNode, we replace the entire statement, not only the NameNode
+        if isinstance(node.expr, NameNode) and node.expr.name in self.substitute:
+            return self.substitute[node.expr.name]
+        else:
+            return self.visit_Node(node)
+    
     def __call__(self, node, substitute):
         self.substitute = substitute
         return super(SubstitutionTransform, self).__call__(node)