A better subpackage redirection

numba/core/utils.py

@@ -12,6 +12,7 @@
 import weakref
 import warnings
 from types import ModuleType
+from importlib import import_module
 from collections.abc import Mapping
 import numpy as np
 
@@ -555,3 +556,38 @@ def unified_function_type(numba_types, require_precise=True):
         function = types.UndefinedFunctionType(mnargs, dispatchers)
 
     return function
+
+
+class _RedirectSubpackage(ModuleType):
+    """Redirect a subpackage to a subpackage.
+
+    This allows all references like:
+
+    >>> from numba.old_subpackage import module
+    >>> module.item
+
+    >>> import numba.old_subpackage.module
+    >>> numba.old_subpackage.module.item
+
+    >>> from numba.old_subpackage.module import item
+    """
+    def __init__(self, old_module_locals, new_module):
+        old_module = old_module_locals['__name__']
+        super().__init__(old_module)
+
+        new_mod_obj = import_module(new_module)
+
+        # Map all sub-modules over
+        for k, v in new_mod_obj.__dict__.items():
+            # Get attributes so that `subpackage.xyz` and
+            # `from subpackage import xyz` work
+            setattr(self, k, v)
+            if isinstance(v, ModuleType):
+                # Map modules into the interpreter so that
+                # `import subpackage.xyz` works
+                sys.modules[f"{old_module}.{k}"] = sys.modules[v.__name__]
+
+        # copy across dunders so that package imports work too
+        for attr, value in old_module_locals.items():
+            if attr.startswith('__') and attr.endswith('__'):
+                setattr(self, attr, value)

numba/tests/test_moved_modules.py

@@ -9,8 +9,8 @@ class TestMovedModule(TestCase):
     def tests_numba_types(self):
         import numba.types
         import numba.core.types as types
-        # The old module is the new module
-        self.assertIs(numba.types, types)
+        # The old module IS NOT the new module
+        self.assertIsNot(numba.types, types)
         # Attribute access are there
         self.assertIs(numba.types.intp, types.intp)
         self.assertIs(numba.types.float64, types.float64)
@@ -19,3 +19,12 @@ def tests_numba_types(self):
         import numba.types.misc
         self.assertIs(types.misc, numba.types.misc)
         self.assertIs(types.misc.Optional, numba.types.misc.Optional)
+        # Import time code could be executed twice and causes the following to
+        # fail.
+        self.assertIs(types.StringLiteral, numba.types.misc.StringLiteral)
+        # Check numba.types.container
+        from numba.types import containers
+        self.assertIs(types.containers, containers)
+        self.assertIs(types.containers.Sequence, containers.Sequence)
+        from numba.types.containers import Sequence
+        self.assertIs(Sequence, containers.Sequence)

numba/types/__init__.py

@@ -1,3 +1,4 @@
 import sys
-import numba.core.types
-sys.modules[__name__] = numba.core.types
+from numba.core.utils import _RedirectSubpackage
+sys.modules[__name__] = _RedirectSubpackage(locals(), "numba.core.types")
+