Fix #611 pkinspect.import_submodule

pykern/pkinspect.py

@@ -184,6 +184,38 @@ def caller_module(exclude_first=True):
     return caller(exclude_first=exclude_first)._module
 
 
+def import_submodule(submodule, subpackage=None, root_packages=None):
+    """Import a module of the form ``root.subpackage.submodule``
+
+    Search ``root_packages``, e.g. ``(sirepo, pykern)``, for modules
+    within the a subpackage of the roots.
+
+    Args:
+        submodule (str): last component of the full module name
+        subpackage (str): name of "middle" component in full module name [submodule_name(caller_module())]
+        packages (iterable): list of packages to search [root_package(caller_module())]
+    Returns:
+        module: imported module object
+    """
+    if root_packages is None:
+        root_packages = (root_package(caller_module(exclude_first=False)),)
+    if subpackage is None:
+        subpackage = submodule_name(caller_module(exclude_first=False))
+    for p in root_packages:
+        s = f"{p}.{subpackage}"
+        n = f"{s}.{submodule}"
+        try:
+            return importlib.import_module(n)
+        except ModuleNotFoundError as e:
+            if e.name not in (p, s, n):
+                # import is failing due to ModuleNotFoundError in a sub-import
+                # not the module we are looking for.
+                raise
+    raise ValueError(
+        f"cannot find module={subpackage}.{submodule} in root_packages={root_packages}"
+    )
+
+
 def is_caller_main():
     """Is the caller's calling module __main__?
 
@@ -239,8 +271,7 @@ def module_name_split(obj):
     Returns:
         str: base part of the module name
     """
-    n = inspect.getmodule(obj).__name__
-    return n.split(".")
+    return inspect.getmodule(obj).__name__.split(".")
 
 
 def module_functions(func_prefix, module=None):

pykern/pkunit.py

@@ -275,6 +275,17 @@ def file_eq(expect_path, *args, **kwargs):
     _FileEq(expect_path, *args, **kwargs)
 
 
+@contextlib.contextmanager
+def insert_data_dir_in_sys_path():
+    """Context manager to insert `data_dir` first in `sys.path`"""
+    p = sys.path
+    try:
+        sys.path = [str(data_dir())] + p
+        yield
+    finally:
+        sys.path = p
+
+
 def is_test_run():
     """Running in a test?
 

tests/pkinspect_data/p1/subpkg1/__init__.py

@@ -0,0 +1 @@
+#

tests/pkinspect_data/p1/subpkg1/err2.py

@@ -0,0 +1 @@
+import error_in_p1

tests/pkinspect_data/p1/subpkg1/mod1.py

@@ -0,0 +1 @@
+#

tests/pkinspect_data/p1/subpkg1/mod3.py

@@ -0,0 +1 @@
+#

tests/pkinspect_data/p2/__init__.py

@@ -0,0 +1 @@
+#

tests/pkinspect_data/p2/subpkg1/__init__.py

@@ -0,0 +1,5 @@
+import pykern.pkinspect
+
+
+def import_submodule(submodule, subpackage, root_packages):
+    return pykern.pkinspect.import_submodule(submodule, subpackage, root_packages)

tests/pkinspect_data/p2/subpkg1/err2.py

@@ -0,0 +1 @@
+import error_in_p2

tests/pkinspect_data/p2/subpkg1/err4.py

@@ -0,0 +1 @@
+import error_in_p2

tests/pkinspect_data/p2/subpkg1/mod1.py

@@ -0,0 +1 @@
+#

tests/pkinspect_data/p2/subpkg1/mod2.py

@@ -0,0 +1 @@
+#

tests/pkinspect_test.py

@@ -1,11 +1,8 @@
-# -*- coding: utf-8 -*-
 """PyTest for :mod:`pykern.pkinspect`
 
 :copyright: Copyright (c) 2015 RadiaSoft LLC.  All Rights Reserved.
 :license: http://www.apache.org/licenses/LICENSE-2.0.html
 """
-from __future__ import absolute_import, division, print_function
-import pytest
 
 
 def test_append_exception_reason():
@@ -84,6 +81,35 @@ def test_caller_module():
     assert expect == n, "{}: should be {}".format(n, expect)
 
 
+def test_import_submodule():
+    from pykern import pkunit, pkinspect, pkcollections
+
+    with pkunit.insert_data_dir_in_sys_path():
+        from p2 import subpkg1
+
+        def _conf(expect, *args):
+            pkcollections.unchecked_del(expect)
+            pkunit.pkeq(expect, subpkg1.import_submodule(*args).__name__)
+
+        _conf("p2.subpkg1.mod1", "mod1", None, None)
+        _conf("p1.subpkg1.mod1", "mod1", None, ("p1", "p2"))
+        _conf("p1.subpkg1.mod1", "mod1", None, ("p1", "p2"))
+        _conf("p2.subpkg1.mod2", "mod2", None, ("p1", "p2"))
+        _conf("p1.subpkg1.mod3", "mod3", None, ("p1", "p2"))
+
+        def _dev(expect, *args):
+            with pkunit.pkexcept(expect):
+                subpkg1.import_submodule(*args)
+
+        _dev("error_in_p2", "err2", None, None)
+        _dev("error_in_p1", "err2", None, ("p1", "p2"))
+        _dev("error_in_p1", "err2", None, ("p1", "p2"))
+        _dev("error_in_p2", "err4", None, ("p1", "p2"))
+        _dev("find module=subpkg1.mod737", "mod737", None, ("p1", "p2"))
+        _dev("find module=subpkg737.mod1", "mod1", "subpkg737", ("p1", "p2"))
+        _dev(r"root_packages=\('not_root_pkg',\)", "mod1", "subpkg1", ("not_root_pkg",))
+
+
 def test_is_caller_main():
     import sys
     import subprocess