FIX: use import system to resolve file executed by kernprof -m (#389)

* Optionally use `importlib` in `kernprof.find_module_script()`

* Added test

* changelog

* Test both explicit profiling and autoprofiling

* Fixes

- `kernprof.py::find_module_script()`:
  Added context manager around call to `importlib.util.find_spec()` to
  minimize side effects
- `tests/test_kernprof.py::test_kernprof_sys_restoration()`:
  Updated test implementation to better protect `sys.modules` (because
  the test is run in-process)

* Don't revert `sys.modules` in `find_module_script()`

Author: TTsangSC

CHANGELOG.rst

@@ -8,6 +8,7 @@ Changes
 * FIX: mitigate speed regressions introduced in 5.0.0
 * ENH: Added capability to combine profiling data both programmatically (``LineStats.__add__()``) and via the CLI (``python -m line_profiler``) (#380, originally proposed in #219)
 * FIX: search function in online documentation
+* FIX: Use import system to locate module file run by ``kernprof -m`` #389
 
 5.0.0
 ~~~~~

kernprof.py

@@ -307,12 +307,26 @@ def stop(self):
         self.is_running = False
 
 
-def find_module_script(module_name, *, exit_on_error=True):
+def find_module_script(module_name, *, static=True, exit_on_error=True):
     """Find the path to the executable script for a module or package."""
     from line_profiler.autoprofile.util_static import modname_to_modpath
+    from importlib.util import find_spec
+
+    def resolve_module_path(mod_name):  # type: (str) -> str | None
+        try:
+            mod_spec = find_spec(mod_name)
+        except ImportError:
+            return None
+        if not mod_spec:
+            return None
+        fname = mod_spec.origin  # type: str | None
+        if fname and os.path.exists(fname):
+            return fname
+
+    get_module_path = modname_to_modpath if static else resolve_module_path
 
     for suffix in '.__main__', '':
-        fname = modname_to_modpath(module_name + suffix)
+        fname = get_module_path(module_name + suffix)
         if fname:
             return fname
 
@@ -1160,8 +1174,8 @@ def _pre_profile(options, module, exit_on_error):
         builtins.__dict__['profile'] = prof
 
     if module:
-        script_file = find_module_script(options.script,
-                                         exit_on_error=exit_on_error)
+        script_file = find_module_script(
+            options.script, static=options.static, exit_on_error=exit_on_error)
     else:
         script_file = find_script(options.script, exit_on_error=exit_on_error)
         # Make sure the script's directory is on sys.path instead of

tests/test_kernprof.py

@@ -125,6 +125,59 @@ def main():
     assert ('Function: main' in proc.stdout) == profiled_main
 
 
+@pytest.mark.parametrize('autoprof', [True, False])
+@pytest.mark.parametrize('static', [True, False])
+def test_kernprof_m_import_resolution(static, autoprof):
+    """
+    Test that `kernprof -m` resolves the module using static and dynamic
+    as is appropriate (toggled by the undocumented environment variable
+    :env:`LINE_PROFILER_STATIC_ANALYSIS`; note that static analysis
+    doesn't handle namespace modules while dynamic resolution does.
+    """
+    code = ub.codeblock('''
+    import enum
+    import os
+    import sys
+
+
+    @profile
+    def main():
+        print('Hello world')
+
+
+    if __name__ == '__main__':
+        main()
+    ''')
+    cmd = [sys.executable, '-m', 'kernprof', '-lv']
+    if autoprof:
+        # Remove the explicit decorator, and use the `--prof-mod` flag
+        code = '\n'.join(line for line in code.splitlines()
+                         if '@profile' not in line)
+        cmd += ['-p', 'my_namesapce_pkg.mysubmod']
+    with tempfile.TemporaryDirectory() as tmpdir:
+        temp_dpath = ub.Path(tmpdir)
+        namespace_mod_path = temp_dpath / 'my_namesapce_pkg' / 'mysubmod.py'
+        namespace_mod_path.parent.mkdir()
+        namespace_mod_path.write_text(code)
+        python_path = tmpdir
+        if 'PYTHONPATH' in os.environ:
+            python_path += ':' + os.environ['PYTHONPATH']
+        env = {**os.environ,
+               # Toggle use of static analysis
+               'LINE_PROFILER_STATIC_ANALYSIS': str(bool(static)),
+               # Add the tempdir to `sys.path`
+               'PYTHONPATH': python_path}
+        cmd += ['-m', 'my_namesapce_pkg.mysubmod']
+        proc = ub.cmd(cmd, cwd=temp_dpath, verbose=2, env=env)
+    if static:
+        assert proc.returncode
+        assert proc.stderr.startswith('Could not find module')
+    else:
+        proc.check_returncode()
+        assert proc.stdout.startswith('Hello world\n')
+        assert 'Function: main' in proc.stdout
+
+
 @pytest.mark.parametrize('error', [True, False])
 @pytest.mark.parametrize(
     'args',
@@ -169,13 +222,18 @@ def main():
             ctx = pytest.raises(BaseException)
         else:
             ctx = contextlib.nullcontext()
-        old_main = sys.modules.get('__main__')
-        with ctx:
-            main(['-l', *shlex.split(args), '-m', 'mymod'])
-        out, _ = capsys.readouterr()
-        assert out.startswith('1')
-        assert tmpdir not in sys.path
-        assert sys.modules.get('__main__') is old_main
+        old_modules = sys.modules.copy()
+        try:
+            old_main = sys.modules.get('__main__')
+            with ctx:
+                main(['-l', *shlex.split(args), '-m', 'mymod'])
+            out, _ = capsys.readouterr()
+            assert out.startswith('1')
+            assert tmpdir not in sys.path
+            assert sys.modules.get('__main__') is old_main
+        finally:
+            sys.modules.clear()
+            sys.modules.update(old_modules)
 
 
 @pytest.mark.parametrize(