Generalize dynamic file coverage to wrap spec_from_file_location()

Instead of wrapping Alembic's load_module_py specifically, wrap
importlib.util.spec_from_file_location() which is the standard way
for tools to dynamically load Python files. This provides coverage
for Alembic and any other tool using this mechanism.

Addresses PR review feedback from @Ar-b-ra.

Co-Authored-By: Claude (global.anthropic.claude-opus-4-5-20251101-v1:0) <noreply@anthropic.com>

Author: emeryberger

src/slipcover/__init__.py

@@ -1,4 +1,4 @@
 from .version import __version__
 from .slipcover import Slipcover, merge_coverage, print_coverage, print_xml
-from .importer import FileMatcher, ImportManager, wrap_pytest, wrap_alembic
+from .importer import FileMatcher, ImportManager, wrap_pytest, wrap_spec_from_file_location
 from .fuzz import wrap_function

src/slipcover/__main__.py

@@ -234,7 +234,7 @@ def main():
     if not args.dont_wrap_pytest:
         sc.wrap_pytest(sci, file_matcher)
 
-    sc.wrap_alembic(sci, file_matcher)
+    sc.wrap_spec_from_file_location(sci, file_matcher)
 
     # Set environment variables for pytest-xdist workers to pick up
     if args.module and args.module[0] == 'pytest':

src/slipcover/importer.py

@@ -281,43 +281,41 @@ def write_pyc(*args, **kwargs):
         pyrewrite.rewrite_asserts = rewrite_asserts_wrapper
 
 
-def wrap_alembic(sci: Slipcover, file_matcher: FileMatcher):
-    """Wraps Alembic's module loading to instrument migration files."""
-    try:
-        import alembic.util.pyfiles as pyfiles
-    except ModuleNotFoundError:
-        return
+def wrap_spec_from_file_location(sci: Slipcover, file_matcher: FileMatcher):
+    """Wraps importlib.util.spec_from_file_location() to instrument dynamically loaded files.
 
+    This provides coverage for tools like Alembic, and any other code that uses
+    spec_from_file_location() to dynamically load Python files.
+    """
     import importlib.util
-    import ast
-    from importlib import machinery
 
-    orig_load_module_py = pyfiles.load_module_py
+    orig_spec_from_file_location = importlib.util.spec_from_file_location
+
+    def spec_from_file_location_wrapper(name, location=None, *, loader=None, submodule_search_locations=None):
+        spec = orig_spec_from_file_location(
+            name, location, loader=loader,
+            submodule_search_locations=submodule_search_locations
+        )
 
-    def load_module_py_wrapper(module_id, path):
-        path = Path(path)
-        if not file_matcher.matches(path):
-            return orig_load_module_py(module_id, path)
+        if spec is None or spec.loader is None:
+            return spec
 
-        # Load and instrument the module
-        spec = importlib.util.spec_from_file_location(module_id, path)
-        assert spec and spec.loader
-        module = importlib.util.module_from_spec(spec)
+        # Skip pytest's assertion rewriting hook - wrap_pytest handles those.
+        # AssertionRewritingHook doesn't have get_code() method.
+        loader_type = type(spec.loader).__name__
+        if loader_type == 'AssertionRewritingHook':
+            return spec
 
-        # Get the code object - handle branch pre-instrumentation if needed
-        if sci.branch and isinstance(spec.loader, machinery.SourceFileLoader) and path.exists():
-            t = br.preinstrument(ast.parse(path.read_bytes()))
-            code = compile(t, str(path), "exec")
-        else:
-            code = spec.loader.get_code(module_id)  # type: ignore[attr-defined]
+        # Skip extension file loaders - can't instrument native extensions
+        if isinstance(spec.loader, machinery.ExtensionFileLoader):
+            return spec
 
-        if code is not None:
-            code = sci.instrument(code)
-            exec(code, module.__dict__)
-        else:
-            # Fallback to original if we can't get the code
-            spec.loader.exec_module(module)  # type: ignore[union-attr]
+        # Check if this file should be instrumented
+        origin = spec.origin or (str(location) if location else None)
+        if origin and file_matcher.matches(origin):
+            # Wrap the loader with our instrumented loader
+            spec.loader = SlipcoverLoader(sci, spec.loader, origin)
 
-        return module
+        return spec
 
-    pyfiles.load_module_py = load_module_py_wrapper
+    importlib.util.spec_from_file_location = spec_from_file_location_wrapper

tests/test_importer.py

@@ -293,8 +293,65 @@ def test_run_script_argv_is_str(tmp_path):
 
 
 @pytest.mark.skipif(sys.platform == 'win32', reason='Fails due to weird PermissionError')
-def test_wrap_alembic(tmp_path, monkeypatch):
-    """Test that Alembic migrations are covered when using wrap_alembic."""
+def test_wrap_spec_from_file_location(tmp_path, monkeypatch):
+    """Test that files loaded via spec_from_file_location are covered."""
+    import json
+
+    # Create a Python file to be loaded dynamically
+    dynamic_module = tmp_path / "dynamic_module.py"
+    dynamic_module.write_text('''
+x = 1  # line 2
+y = 2  # line 3
+z = x + y  # line 4
+''')
+
+    # Create a script that loads the module via spec_from_file_location
+    script = tmp_path / "main_script.py"
+    script.write_text(f"""
+import importlib.util
+spec = importlib.util.spec_from_file_location("dynamic_module", "{dynamic_module}")
+module = importlib.util.module_from_spec(spec)
+spec.loader.exec_module(module)
+assert module.z == 3
+""")
+
+    monkeypatch.chdir(tmp_path)
+
+    out = tmp_path / "coverage.json"
+    result = subprocess.run(
+        [sys.executable, "-m", "slipcover", "--json", "--out", str(out),
+         "--source", str(tmp_path), str(script)],
+        capture_output=True,
+        text=True
+    )
+
+    assert result.returncode == 0, f"stdout: {result.stdout}\nstderr: {result.stderr}"
+
+    with out.open() as f:
+        cov = json.load(f)
+
+    # Check that the dynamically loaded file was covered
+    dynamic_files = [k for k in cov['files'].keys() if 'dynamic_module.py' in k]
+    assert dynamic_files, f"Dynamic module not in coverage: {list(cov['files'].keys())}"
+
+    # Check that lines were executed
+    file_cov = cov['files'][dynamic_files[0]]
+    executed_lines = file_cov['executed_lines']
+    assert 2 in executed_lines and 3 in executed_lines and 4 in executed_lines, \
+        f"Lines not executed, got: {executed_lines}"
+
+
+try:
+    import alembic
+    HAS_ALEMBIC = True
+except ImportError:
+    HAS_ALEMBIC = False
+
+
+@pytest.mark.skipif(sys.platform == 'win32', reason='Fails due to weird PermissionError')
+@pytest.mark.skipif(not HAS_ALEMBIC, reason='Alembic not installed')
+def test_wrap_spec_from_file_location_alembic(tmp_path, monkeypatch):
+    """Test that Alembic migrations are covered (integration test for spec_from_file_location)."""
     import json
 
     # Create a minimal alembic setup
@@ -394,8 +451,59 @@ def downgrade():
 
 
 @pytest.mark.skipif(sys.platform == 'win32', reason='Fails due to weird PermissionError')
-def test_wrap_alembic_with_branch(tmp_path, monkeypatch):
-    """Test that Alembic migrations are covered with branch coverage enabled."""
+def test_wrap_spec_from_file_location_with_branch(tmp_path, monkeypatch):
+    """Test that files loaded via spec_from_file_location get branch coverage."""
+    import json
+
+    # Create a Python file with a branch to be loaded dynamically
+    dynamic_module = tmp_path / "dynamic_module.py"
+    dynamic_module.write_text('''
+x = 1
+if x > 0:  # branch
+    y = 2
+else:
+    y = 3
+z = y
+''')
+
+    # Create a script that loads the module via spec_from_file_location
+    script = tmp_path / "main_script.py"
+    script.write_text(f"""
+import importlib.util
+spec = importlib.util.spec_from_file_location("dynamic_module", "{dynamic_module}")
+module = importlib.util.module_from_spec(spec)
+spec.loader.exec_module(module)
+assert module.z == 2
+""")
+
+    monkeypatch.chdir(tmp_path)
+
+    out = tmp_path / "coverage.json"
+    result = subprocess.run(
+        [sys.executable, "-m", "slipcover", "--branch", "--json", "--out", str(out),
+         "--source", str(tmp_path), str(script)],
+        capture_output=True,
+        text=True
+    )
+
+    assert result.returncode == 0, f"stdout: {result.stdout}\nstderr: {result.stderr}"
+
+    with out.open() as f:
+        cov = json.load(f)
+
+    # Check that the dynamically loaded file was covered
+    dynamic_files = [k for k in cov['files'].keys() if 'dynamic_module.py' in k]
+    assert dynamic_files, f"Dynamic module not in coverage: {list(cov['files'].keys())}"
+
+    # Check that branch info is present
+    file_cov = cov['files'][dynamic_files[0]]
+    assert 'executed_branches' in file_cov, "Branch coverage not recorded"
+
+
+@pytest.mark.skipif(sys.platform == 'win32', reason='Fails due to weird PermissionError')
+@pytest.mark.skipif(not HAS_ALEMBIC, reason='Alembic not installed')
+def test_wrap_spec_from_file_location_with_branch_alembic(tmp_path, monkeypatch):
+    """Test that Alembic migrations get branch coverage (integration test)."""
     import json
 
     # Create a minimal alembic setup