Make TreeFragment handle DefNode and PropertyNode names as well

Cython/Compiler/Tests/TestTreeFragment.py

@@ -162,6 +162,22 @@ class TestTreeFragments(CythonTest):
         self.assertTrue(T.stats[0].expr.obj.name == "x", T)
         self.assertTrue(T.stats[0].expr.attribute == "y", T)
 
+    def test_defnode(self):
+        F = self.fragment(u"def NAME(): pass")
+        T = F.substitute({
+            "NAME" : "test",
+        })
+        self.assertTrue(isinstance(T.stats[0], DefNode), T)
+        self.assertTrue(T.stats[0].name == "test", T)
+
+    def test_propertynode(self):
+        F = TreeFragment(u"property NAME: pass", level='c_class')
+        T = F.substitute({
+            "NAME" : "test",
+        })
+        self.assertTrue(isinstance(T.stats[0], PropertyNode), T)
+        self.assertTrue(T.stats[0].name == "test", T)
+
 if __name__ == "__main__":
     import unittest
     unittest.main()

Cython/Compiler/TreeFragment.py

@@ -206,13 +206,16 @@ class TemplateTransform(VisitorTransform):
             ret.declarator = CNameDeclaratorNode(pos=node.pos, name=EncodedString(""), cname=None)
         return ret
 
-    def visit_AttributeNode(self, node):
-        ret = self.visit_Node(node)
-        sub = self.substitutions.get(node.attribute)
+    def try_substitute_string_attribute(self, node, attrname):
+        sub = self.substitutions.get(getattr(node, attrname))
         if isinstance(sub, EncodedString):
-            ret.attribute = sub
+            setattr(node, attrname, sub)
         elif isinstance(sub, str):
-            ret.attribute = EncodedString(sub)
+            setattr(node, attrname, EncodedString(sub))
+
+    def visit_AttributeNode(self, node):
+        ret = self.visit_Node(node)
+        self.try_substitute_string_attribute(ret, "attribute")
         return ret
 
     def visit_CArgDeclNode(self, node):
@@ -222,6 +225,16 @@ class TemplateTransform(VisitorTransform):
                 return sub
         return self.visit_Node(node)
 
+    def visit_DefNode(self, node):
+        ret = self.visit_Node(node)
+        self.try_substitute_string_attribute(ret, "name")
+        return ret
+
+    def visit_PropertyNode(self, node):
+        ret = self.visit_Node(node)
+        self.try_substitute_string_attribute(ret, "name")
+        return ret
+
     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