[transforms] allow transformation functions on any nodes, not only modules

6 files changed, 47 insertions, 14 deletions

diff --git a/ChangeLog b/ChangeLog
index 6d5e5f89..e8df514c 100644
--- a/ChangeLog
+++ b/ChangeLog
@@ -2,8 +2,13 @@ Change log for the astroid package (used to be astng)
 =====================================================
 
 --
+    * Allow transformation functions on any node, providing a
+      `register_transform` function on the manager instead of the 
+     `register_transformer` to make it more flexible wrt node selection
+
     * Added the test_utils module for building ASTs and
       extracting deeply nested nodes for easier testing.
+
     * rename the project as astroid
 
 2013-04-16  --  0.24.3
diff --git a/brain/py2stdlib.py b/brain/py2stdlib.py
index 25c7122d..b3da2a53 100644
--- a/brain/py2stdlib.py
+++ b/brain/py2stdlib.py
@@ -5,11 +5,14 @@ Currently help understanding of :
 * hashlib.md5 and hashlib.sha1
 """
 
-from astroid import MANAGER
+from astroid import MANAGER, nodes
 from astroid.builder import AstroidBuilder
 
 MODULE_TRANSFORMS = {}
 
+
+# module specific transformation functions #####################################
+
 def hashlib_transform(module):
     fake = AstroidBuilder(MANAGER).string_build('''
 
@@ -177,6 +180,6 @@ def transform(module):
         tr(module)
 
 from astroid import MANAGER
-MANAGER.register_transformer(transform)
+MANAGER.register_transform(nodes.Module, transform)
 
 
diff --git a/builder.py b/builder.py
index 42168eb5..6b80f350 100644
--- a/builder.py
+++ b/builder.py
@@ -138,9 +138,6 @@ class AstroidBuilder(InspectBuilder):
         # handle delayed assattr nodes
         for delayed in module._delayed_assattr:
             self.delayed_assattr(delayed)
-        if modname:
-            for transformer in self._manager.transformers:
-                transformer(module)
         return module
 
     def _data_build(self, data, modname, path):
@@ -156,7 +153,7 @@ class AstroidBuilder(InspectBuilder):
             package = True
         else:
             package = path and path.find('__init__.py') > -1 or False
-        rebuilder = TreeRebuilder()
+        rebuilder = TreeRebuilder(self._manager)
         module = rebuilder.visit_module(node, modname, package)
         module.file = module.path = node_file
         module._from_nodes = rebuilder._from_nodes
diff --git a/manager.py b/manager.py
index 18b48d8e..1e3091c4 100644
--- a/manager.py
+++ b/manager.py
@@ -83,7 +83,7 @@ class AstroidManager(OptionsProviderMixIn):
             # NOTE: cache entries are added by the [re]builder
             self.astroid_cache = {}
             self._mod_file_cache = {}
-            self.transformers = []
+            self.transforms = {}
 
     def astroid_from_file(self, filepath, modname=None, fallback=True, source=False):
         """given a module name, return the astroid object"""
@@ -263,8 +263,15 @@ class AstroidManager(OptionsProviderMixIn):
                     project.add_module(astroid)
         return project
 
-    def register_transformer(self, transformer):
-        self.transformers.append(transformer)
+    def register_transform(self, node_class, transform, predicate=None):
+        """Register `transform(node)` function to be applied on the given
+        Astroid's `node_class` if `predicate` is None or return a true value
+        when called with the node as argument.
+
+        The transform function may return a value which is then used to
+        substitute the original node in the tree.
+        """
+        self.transforms.setdefault(node_class, []).append( (transform, predicate) )
 
 class Project:
     """a project handle a set of modules / packages"""
diff --git a/rebuilder.py b/rebuilder.py
index 2ce6a19f..47c81138 100644
--- a/rebuilder.py
+++ b/rebuilder.py
@@ -20,6 +20,7 @@ order to get a single Astroid representation
 """
 
 import sys
+from warnings import warn
 from _ast import (Expr as Discard, Str,
     # binary operators
     Add, Div, FloorDiv,  Mod, Mult, Pow, Sub, BitAnd, BitOr, BitXor,
@@ -119,7 +120,8 @@ def _set_infos(oldnode, newnode, parent):
 class TreeRebuilder(object):
     """Rebuilds the _ast tree to become an Astroid tree"""
 
-    def __init__(self):
+    def __init__(self, manager):
+        self._manager = manager
         self.asscontext = None
         self._metaclass = ['']
         self._global_names = []
@@ -127,6 +129,25 @@ class TreeRebuilder(object):
         self._delayed_assattr = []
         self._visit_meths = {}
 
+    def _transform(self, node):
+        try:
+            transforms = self._manager.transforms[type(node)]
+        except KeyError:
+            return node # no transform registered for this class of node
+        orig_node = node # copy the reference
+        for transform_func, predicate in transforms:
+            if predicate is not None and predicate(node):
+                ret = transform_func(node)
+                # if the transformation function returns something, it's
+                # expected to be a replacement for the node
+                if ret is not None:
+                    if node is not orig_node:
+                        # node has already be modified by some previous
+                        # transformation, warn about it
+                        warn('node %s substitued multiple times' % node)
+                    node = ret
+        return node
+
     def visit_module(self, node, modname, package):
         """visit a Module node by returning a fresh instance of it"""
         newnode = new.Module(modname, None)
@@ -135,18 +156,18 @@ class TreeRebuilder(object):
         _init_set_doc(node, newnode)
         newnode.body = [self.visit(child, newnode) for child in node.body]
         newnode.set_line_info(newnode.last_child())
-        return newnode
+        return self._transform(newnode)
 
     def visit(self, node, parent):
         cls = node.__class__
         if cls in self._visit_meths:
-            return self._visit_meths[cls](node, parent)
+            visit_method = self._visit_meths[cls]
         else:
             cls_name = cls.__name__
             visit_name = 'visit_' + REDIRECT.get(cls_name, cls_name).lower()
             visit_method = getattr(self, visit_name)
             self._visit_meths[cls] = visit_method
-            return visit_method(node, parent)
+        return self._transform(visit_method(node, parent))
 
     def _save_assignment(self, node, name=None):
         """save assignement situation since node.parent is not available yet"""
diff --git a/test/unittest_regrtest.py b/test/unittest_regrtest.py
index 5cc114e5..b5878bc6 100644
--- a/test/unittest_regrtest.py
+++ b/test/unittest_regrtest.py
@@ -41,7 +41,7 @@ class NonRegressionTC(TestCase):
         manager.__dict__ = {}
         manager.astroid_cache = {}
         manager._mod_file_cache = {}
-        manager.transformers = {}
+        manager.transforms = {}
         return manager
 
     def test_module_path(self):