Simplify TrackNumpyAttributes transform.

1 files changed, 8 insertions, 3 deletions

diff --git a/Cython/Compiler/ParseTreeTransforms.py b/Cython/Compiler/ParseTreeTransforms.py
index aee77f81d..9e720c941 100644
--- a/Cython/Compiler/ParseTreeTransforms.py
+++ b/Cython/Compiler/ParseTreeTransforms.py
@@ -611,9 +611,11 @@ class PxdPostParse(CythonTransform, SkipDeclarations):
         else:
             return node
 
-class TrackNumpyAttributes(CythonTransform, SkipDeclarations):
-    def __init__(self, context):
-        super(TrackNumpyAttributes, self).__init__(context)
+
+class TrackNumpyAttributes(VisitorTransform, SkipDeclarations):
+    # TODO: Make name handling as good as in InterpretCompilerDirectives() below - probably best to merge the two.
+    def __init__(self):
+        super(TrackNumpyAttributes, self).__init__()
         self.numpy_module_names = set()
 
     def visit_CImportStatNode(self, node):
@@ -627,6 +629,9 @@ class TrackNumpyAttributes(CythonTransform, SkipDeclarations):
             node.is_numpy_attribute = True
         return node
 
+    visit_Node = VisitorTransform.recurse_to_children
+
+
 class InterpretCompilerDirectives(CythonTransform, SkipDeclarations):
     """
     After parsing, directives can be stored in a number of places: