diff options
| -rw-r--r-- | builder.py | 6 | ||||
| -rw-r--r-- | manager.py | 27 | ||||
| -rw-r--r-- | rebuilder.py | 20 | ||||
| -rw-r--r-- | test/unittest_builder.py | 13 |
4 files changed, 47 insertions, 19 deletions
@@ -102,6 +102,12 @@ class AstroidBuilder(InspectBuilder): # this is a built-in module # get a partial representation by introspection node = self.inspect_build(module, modname=modname, path=path) + # we have to handle transformation by ourselves since the rebuilder + # isn't called for builtin nodes + # + # XXX it's then only called for Module nodes, not for underlying + # nodes + node = self._manager.transform(node) return node def file_build(self, path, modname=None): @@ -273,6 +273,33 @@ class AstroidManager(OptionsProviderMixIn): """ self.transforms.setdefault(node_class, []).append( (transform, predicate) ) + def unregister_transform(self, node_class, transform, predicate=None): + """Unregister the given transform.""" + self.transforms[node_class].remove( (transform, predicate) ) + + def transform(self, node): + """Call matching transforms for the given node if any and return the + transformed node. + """ + try: + transforms = self.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 None or 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 + + class Project(object): """a project handle a set of modules / packages""" diff --git a/rebuilder.py b/rebuilder.py index f6fae7b1..ef8e7635 100644 --- a/rebuilder.py +++ b/rebuilder.py @@ -142,25 +142,7 @@ class TreeRebuilder(object): self._from_nodes = [] 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 None or 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 + self._transform = manager.transform def visit_module(self, node, modname, package): """visit a Module node by returning a fresh instance of it""" diff --git a/test/unittest_builder.py b/test/unittest_builder.py index 9b29852b..54115393 100644 --- a/test/unittest_builder.py +++ b/test/unittest_builder.py @@ -350,6 +350,19 @@ class BuilderTC(TestCase): self.assertEqual(infered.name, 'type') infered.as_string() # no crash test + def test_inspect_transform_module(self): + # ensure no cached version of the time module + MANAGER._mod_file_cache.pop(('time', None), None) + def transform_time(node): + if node.name == 'time': + node.transformed = True + MANAGER.register_transform(nodes.Module, transform_time) + try: + time_ast = MANAGER.ast_from_module_name('time') + self.assertTrue(getattr(time_ast, 'transformed', False)) + finally: + MANAGER.unregister_transform(nodes.Module, transform_time) + def test_package_name(self): """test base properties and method of a astroid module""" datap = self.builder.file_build(join(DATA, '__init__.py'), 'data') |