Add coverage support for Alembic migrations (Issue #76)

Alembic loads migration files using importlib.util.spec_from_file_location()
which bypasses Slipcover's sys.meta_path import hooks. This adds wrap_alembic()
to intercept Alembic's load_module_py function and instrument migration code
before execution, similar to the existing wrap_pytest() approach.

Co-Authored-By: Claude Opus 4.5 <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
+from .importer import FileMatcher, ImportManager, wrap_pytest, wrap_alembic
 from .fuzz import wrap_function

src/slipcover/__main__.py

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

src/slipcover/importer.py

@@ -279,3 +279,45 @@ def write_pyc(*args, **kwargs):
         pyrewrite._read_pyc = read_pyc
         pyrewrite._write_pyc = write_pyc
         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
+
+    import importlib.util
+    import ast
+    from importlib import machinery
+
+    orig_load_module_py = pyfiles.load_module_py
+
+    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)
+
+        # 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)
+
+        # 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]
+
+        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]
+
+        return module
+
+    pyfiles.load_module_py = load_module_py_wrapper

tests/test_importer.py

@@ -290,3 +290,205 @@ def test_run_script_argv_is_str(tmp_path):
 """)
 
     subprocess.run([sys.executable, "-m", "slipcover", "--silent", cmdfile], check=True)
+
+
+@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."""
+    import json
+
+    # Create a minimal alembic setup
+    migrations_dir = tmp_path / "migrations"
+    versions_dir = migrations_dir / "versions"
+    versions_dir.mkdir(parents=True)
+
+    # Create alembic.ini
+    alembic_ini = tmp_path / "alembic.ini"
+    alembic_ini.write_text(f"""
+[alembic]
+script_location = {migrations_dir}
+sqlalchemy.url = sqlite:///:memory:
+""")
+
+    # Create env.py
+    env_py = migrations_dir / "env.py"
+    env_py.write_text("""
+from alembic import context
+
+def run_migrations_offline():
+    context.configure(url="sqlite:///:memory:", literal_binds=True)
+    with context.begin_transaction():
+        context.run_migrations()
+
+def run_migrations_online():
+    from sqlalchemy import create_engine
+    connectable = create_engine("sqlite:///:memory:")
+    with connectable.connect() as connection:
+        context.configure(connection=connection)
+        with context.begin_transaction():
+            context.run_migrations()
+
+if context.is_offline_mode():
+    run_migrations_offline()
+else:
+    run_migrations_online()
+""")
+
+    # Create script.py.mako (required by alembic)
+    script_mako = migrations_dir / "script.py.mako"
+    script_mako.write_text("")
+
+    # Create a migration file
+    migration_file = versions_dir / "001_test_migration.py"
+    migration_file.write_text('''
+"""test migration"""
+revision = '001'
+down_revision = None
+
+def upgrade():
+    x = 1  # line 8
+    y = 2  # line 9
+
+def downgrade():
+    pass  # line 12
+''')
+
+    # Create a script that runs the alembic migration
+    script = tmp_path / "run_migration.py"
+    script.write_text(f"""
+import sys
+sys.path.insert(0, '{tmp_path}')
+from alembic.config import Config
+from alembic import command
+
+alembic_cfg = Config('{alembic_ini}')
+command.upgrade(alembic_cfg, 'head')
+""")
+
+    monkeypatch.chdir(tmp_path)
+
+    out = tmp_path / "coverage.json"
+    result = subprocess.run(
+        [sys.executable, "-m", "slipcover", "--json", "--out", str(out),
+         "--source", str(versions_dir), str(script)],
+        capture_output=True,
+        text=True
+    )
+
+    # The script should complete (may have warnings, but not crash)
+    assert result.returncode == 0, f"stdout: {result.stdout}\nstderr: {result.stderr}"
+
+    with out.open() as f:
+        cov = json.load(f)
+
+    # Check that the migration file was covered
+    # The path might be stored as relative or absolute depending on the working directory
+    migration_files = [k for k in cov['files'].keys() if '001_test_migration.py' in k]
+    assert migration_files, f"Migration file not in coverage: {list(cov['files'].keys())}"
+
+    # Check that lines in the upgrade function were executed
+    file_cov = cov['files'][migration_files[0]]
+    executed_lines = file_cov['executed_lines']
+    # Lines 7 and 8 are inside upgrade() function (x=1, y=2)
+    assert 7 in executed_lines or 8 in executed_lines, f"upgrade() lines not executed, executed: {executed_lines}"
+
+
+@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."""
+    import json
+
+    # Create a minimal alembic setup
+    migrations_dir = tmp_path / "migrations"
+    versions_dir = migrations_dir / "versions"
+    versions_dir.mkdir(parents=True)
+
+    # Create alembic.ini
+    alembic_ini = tmp_path / "alembic.ini"
+    alembic_ini.write_text(f"""
+[alembic]
+script_location = {migrations_dir}
+sqlalchemy.url = sqlite:///:memory:
+""")
+
+    # Create env.py
+    env_py = migrations_dir / "env.py"
+    env_py.write_text("""
+from alembic import context
+
+def run_migrations_offline():
+    context.configure(url="sqlite:///:memory:", literal_binds=True)
+    with context.begin_transaction():
+        context.run_migrations()
+
+def run_migrations_online():
+    from sqlalchemy import create_engine
+    connectable = create_engine("sqlite:///:memory:")
+    with connectable.connect() as connection:
+        context.configure(connection=connection)
+        with context.begin_transaction():
+            context.run_migrations()
+
+if context.is_offline_mode():
+    run_migrations_offline()
+else:
+    run_migrations_online()
+""")
+
+    # Create script.py.mako (required by alembic)
+    script_mako = migrations_dir / "script.py.mako"
+    script_mako.write_text("")
+
+    # Create a migration file with a branch
+    migration_file = versions_dir / "001_test_migration.py"
+    migration_file.write_text('''
+"""test migration"""
+revision = '001'
+down_revision = None
+
+def upgrade():
+    x = 1
+    if x > 0:  # branch
+        y = 2
+    else:
+        y = 3
+
+def downgrade():
+    pass
+''')
+
+    # Create a script that runs the alembic migration
+    script = tmp_path / "run_migration.py"
+    script.write_text(f"""
+import sys
+sys.path.insert(0, '{tmp_path}')
+from alembic.config import Config
+from alembic import command
+
+alembic_cfg = Config('{alembic_ini}')
+command.upgrade(alembic_cfg, 'head')
+""")
+
+    monkeypatch.chdir(tmp_path)
+
+    out = tmp_path / "coverage.json"
+    result = subprocess.run(
+        [sys.executable, "-m", "slipcover", "--branch", "--json", "--out", str(out),
+         "--source", str(versions_dir), str(script)],
+        capture_output=True,
+        text=True
+    )
+
+    # The script should complete (may have warnings, but not crash)
+    assert result.returncode == 0, f"stdout: {result.stdout}\nstderr: {result.stderr}"
+
+    with out.open() as f:
+        cov = json.load(f)
+
+    # Check that the migration file was covered
+    migration_files = [k for k in cov['files'].keys() if '001_test_migration.py' in k]
+    assert migration_files, f"Migration file not in coverage: {list(cov['files'].keys())}"
+
+    # Check that branch info is present
+    file_cov = cov['files'][migration_files[0]]
+    assert 'executed_branches' in file_cov, "Branch coverage not recorded"