Handle import cycles.

3 files changed, 30 insertions, 3 deletions

diff --git a/NEWS b/NEWS
index 60713b8..ab593b6 100644
--- a/NEWS
+++ b/NEWS
@@ -6,6 +6,9 @@ Changes and improvements to extras_, grouped by release.
 NEXT
 ~~~~
 
+* Imports in the middle of import cycles are now supported.
+  (Robert Collins)
+
 0.0.3
 ~~~~~
 
diff --git a/extras/__init__.py b/extras/__init__.py
index 2d34b52..29d2230 100644
--- a/extras/__init__.py
+++ b/extras/__init__.py
@@ -40,13 +40,17 @@ def try_import(name, alternative=None, error_callback=None):
     """
     module_segments = name.split('.')
     last_error = None
+    remainder = []
+    # module_name will be what successfully imports. We cannot walk from the
+    # __import__ result because in import loops (A imports A.B, which imports
+    # C, which calls try_import("A.B")) A.B will not yet be set.
     while module_segments:
         module_name = '.'.join(module_segments)
         try:
-            module = __import__(module_name)
+            __import__(module_name)
         except ImportError:
             last_error = sys.exc_info()[1]
-            module_segments.pop()
+            remainder.append(module_segments.pop())
             continue
         else:
             break
@@ -54,8 +58,9 @@ def try_import(name, alternative=None, error_callback=None):
         if last_error is not None and error_callback is not None:
             error_callback(last_error)
         return alternative
+    module = sys.modules[module_name]
     nonexistent = object()
-    for segment in name.split('.')[1:]:
+    for segment in reversed(remainder):
         module = getattr(module, segment, nonexistent)
         if module is nonexistent:
             if last_error is not None and error_callback is not None:
diff --git a/extras/tests/test_extras.py b/extras/tests/test_extras.py
index be1ed1c..d048bda 100644
--- a/extras/tests/test_extras.py
+++ b/extras/tests/test_extras.py
@@ -1,5 +1,8 @@
 # Copyright (c) 2010-2012 extras developers. See LICENSE for details.
 
+import sys
+import types
+
 from testtools import TestCase
 from testtools.matchers import (
     Equals,
@@ -125,6 +128,22 @@ class TestTryImport(TestCase):
         # the error callback is not called on success.
         check_error_callback(self, try_import, 'os.path', 0, True)
 
+    def test_handle_partly_imported_name(self):
+        # try_import('thing.other') when thing.other is mid-import
+        # used to fail because thing.other is not assigned until thing.other
+        # finishes its import - but thing.other is accessible via sys.modules.
+        outer = types.ModuleType("extras.outer")
+        inner = types.ModuleType("extras.outer.inner")
+        inner.attribute = object()
+        self.addCleanup(sys.modules.pop, "extras.outer", None)
+        self.addCleanup(sys.modules.pop, "extras.outer.inner", None)
+        sys.modules["extras.outer"] = outer
+        sys.modules["extras.outer.inner"] = inner
+        result = try_import("extras.outer.inner.attribute")
+        self.expectThat(result, Is(inner.attribute))
+        result = try_import("extras.outer.inner")
+        self.expectThat(result, Is(inner))
+
 
 class TestTryImports(TestCase):