Generalize hooks for resolving imports that cannot be resolved normally,

and use the functionality in brain/py2gi.

3 files changed, 40 insertions, 24 deletions

diff --git a/astroid/brain/py2gi.py b/astroid/brain/py2gi.py
index 20356c1..6747898 100644
--- a/astroid/brain/py2gi.py
+++ b/astroid/brain/py2gi.py
@@ -112,26 +112,10 @@ def _gi_build_stub(parent):
 
     return ret
 
-# Overwrite Module.module_import to _actually_ import the introspected module if
-# it's a gi module, then build stub code by examining its info and get an astng
-# from that
-
-from astroid.scoped_nodes import Module
-_orig_import_module = Module.import_module
-
-def _new_import_module(self, modname, relative_only=False, level=None):
-    # Could be a static piece of gi.repository or whatever unrelated module,
-    # let that fall through
-    try:
-        return _orig_import_module(self, modname, relative_only, level)
-    except AstroidBuildingException:
-        # we only consider gi.repository submodules
-        if not modname.startswith('gi.repository.'):
-            if relative_only and level is None:
-                level = 0
-            modname = self.relative_to_absolute_name(modname, level)
-        if not modname.startswith('gi.repository.'):
-            raise
+def _import_gi_module(modname):
+    # we only consider gi.repository submodules
+    if not modname.startswith('gi.repository.'):
+        raise AstroidBuildingException()
     # build astroid representation unless we already tried so
     if modname not in _inspected_modules:
         modnames = [modname]
@@ -166,4 +150,6 @@ def _new_import_module(self, modname, relative_only=False, level=None):
         raise AstroidBuildingException('Failed to import module %r' % modname)
     return astng
 
-Module.import_module = _new_import_module
+
+MANAGER.register_failed_import_hook(_import_gi_module)
+
diff --git a/astroid/manager.py b/astroid/manager.py
index 2f5b555..1e6d36d 100644
--- a/astroid/manager.py
+++ b/astroid/manager.py
@@ -88,6 +88,7 @@ class AstroidManager(OptionsProviderMixIn):
             self.astroid_cache = {}
             self._mod_file_cache = {}
             self.transforms = collections.defaultdict(list)
+            self._failed_import_hooks = []
 
     def ast_from_file(self, filepath, modname=None, fallback=True, source=False):
         """given a module name, return the astroid object"""
@@ -144,6 +145,13 @@ class AstroidManager(OptionsProviderMixIn):
             if filepath is None:
                 raise AstroidBuildingException("Unable to load module %s" % (modname,))
             return self.ast_from_file(filepath, modname, fallback=False)
+        except AstroidBuildingException as e:
+            for hook in self._failed_import_hooks:
+                try:
+                    return hook(modname)
+                except AstroidBuildingException:
+                    pass
+            raise e
         finally:
             os.chdir(old_cwd)
 
@@ -290,6 +298,16 @@ class AstroidManager(OptionsProviderMixIn):
         """Unregister the given transform."""
         self.transforms[node_class].remove((transform, predicate))
 
+    def register_failed_import_hook(self, hook):
+        """"Registers a hook to resolve imports that cannot be found otherwise.
+
+        `hook` must be a function that accepts a single argument `modname` which
+        contains the name of the module or package that could not be imported.
+        If `hook` can resolve the import, must return a node of type `astroid.Module`,
+        otherwise, it must raise `AstroidBuildingException`.
+        """
+        self._failed_import_hooks.append(hook)
+
     def transform(self, node):
         """Call matching transforms for the given node if any and return the
         transformed node.
diff --git a/astroid/tests/unittest_manager.py b/astroid/tests/unittest_manager.py
index 6fdc04d..4ed42e4 100644
--- a/astroid/tests/unittest_manager.py
+++ b/astroid/tests/unittest_manager.py
@@ -124,7 +124,6 @@ class AstroidManagerTest(resources.SysPathSetup, unittest.TestCase):
 
     def test_file_from_module(self):
         """check if the unittest filepath is equals to the result of the method"""
-        import unittest
         if sys.version_info > (3, 0):
             unittest_file = unittest.__file__
         else:
@@ -137,7 +136,6 @@ class AstroidManagerTest(resources.SysPathSetup, unittest.TestCase):
         self.assertRaises(AstroidBuildingException, self.manager.file_from_module_name, 'unhandledModule', None)
 
     def test_ast_from_module(self):
-        import unittest
         astroid = self.manager.ast_from_module(unittest)
         self.assertEqual(astroid.pure_python, True)
         import time
@@ -146,7 +144,6 @@ class AstroidManagerTest(resources.SysPathSetup, unittest.TestCase):
 
     def test_ast_from_module_cache(self):
         """check if the module is in the cache manager"""
-        import unittest
         astroid = self.manager.ast_from_module(unittest)
         self.assertEqual(astroid.name, 'unittest')
         self.assertIn('unittest', self.manager.astroid_cache)
@@ -224,6 +221,21 @@ class AstroidManagerTest(resources.SysPathSetup, unittest.TestCase):
             'data.unicode_package.core']
         self.assertListEqual(sorted(obj.keys()), expected)
 
+    def testFailedImportHooks(self):
+        def hook(modname):
+            if modname == 'foo.bar':
+                return unittest
+            else:
+                raise AstroidBuildingException()
+
+        with self.assertRaises(AstroidBuildingException):
+            self.manager.ast_from_module_name('foo.bar')
+        self.manager.register_failed_import_hook(hook)
+        self.assertEqual(unittest, self.manager.ast_from_module_name('foo.bar'))
+        with self.assertRaises(AstroidBuildingException):
+            self.manager.ast_from_module_name('foo.bar.baz')
+        del self.manager._failed_import_hooks[0]
+
 
 class BorgAstroidManagerTC(unittest.TestCase):