- replaced bytecode editing-based pytest wrapping with an AST-based m…

…ethod,

  so that bytecode editing does not need to be ported to Python 3.12;

- enabled pytest wrapping under Python 3.12;

- removed code no longer needed;

src/slipcover/__main__.py

@@ -36,8 +36,7 @@
 ap.add_argument('--silent', action='store_true', help=argparse.SUPPRESS)
 ap.add_argument('--dis', action='store_true', help=argparse.SUPPRESS)
 ap.add_argument('--debug', action='store_true', help=argparse.SUPPRESS)
-if sys.version_info[0:2] < (3,12):
-    ap.add_argument('--dont-wrap-pytest', action='store_true', help=argparse.SUPPRESS)
+ap.add_argument('--dont-wrap-pytest', action='store_true', help=argparse.SUPPRESS)
 
 g = ap.add_mutually_exclusive_group(required=True)
 g.add_argument('-m', dest='module', nargs=1, help="run given module as __main__")
@@ -73,7 +72,7 @@
                    skip_covered=args.skip_covered, disassemble=args.dis)
 
 
-if sys.version_info[0:2] < (3,12) and not args.dont_wrap_pytest:
+if not args.dont_wrap_pytest:
     sc.wrap_pytest(sci, file_matcher)
 
 

src/slipcover/bytecode.py

@@ -592,56 +592,6 @@ def disable_inserted_function(self, offset):
         self.patch[offset] = op_JUMP_FORWARD
 
 
-    def replace_global_with_const(self, global_name, const_index):
-        """Replaces a global name lookup by a constant load."""
-        assert not self.finished
-
-        if self.patch is None:
-            self.patch = bytearray(self.orig_code.co_code)
-
-        if self.branches is None:
-            self.branches = Branch.from_code(self.orig_code)
-            self.ex_table = ExceptionTableEntry.from_code(self.orig_code)
-            self.lines = LineEntry.from_code(self.orig_code)
-
-        if global_name in self.orig_code.co_names:
-            name_index = self.orig_code.co_names.index(global_name)
-
-            def find_load_globals():
-                for op_off, op_len, op, op_arg in unpack_opargs(self.patch):
-                    if op == op_LOAD_GLOBAL:
-                        if PYTHON_VERSION >= (3,11):
-                            if (op_arg>>1) == name_index:
-                                yield (op_off, op_len, op, op_arg)
-                        else:
-                            if op_arg == name_index:
-                                yield (op_off, op_len, op, op_arg)
-
-            delta = 0
-            # read from pre-computed list() below so we can modify on the fly
-            for op_off, op_len, op, op_arg in list(find_load_globals()):
-                repl = bytearray()
-                if sys.version_info[0:2] >= (3,11) and op_arg&1:
-                    repl.extend(opcode_arg(dis.opmap['PUSH_NULL'], 0))
-                repl.extend(opcode_arg(op_LOAD_CONST, const_index))
-
-                op_off += delta     # adjust for any other changes
-                self.patch[op_off:op_off+op_len] = repl
-
-                change = len(repl) - op_len
-                if change:
-                    for l in self.lines:
-                        l.adjust(op_off, change)
-
-                    for b in self.branches:
-                        b.adjust(op_off, change)
-
-                    for e in self.ex_table:
-                        e.adjust(op_off, change)
-
-                delta += change
-
-
     def _finish(self):
         if not self.finished:
             self.finished = True

src/slipcover/importer.py

@@ -157,24 +157,59 @@ def __exit__(self, *args: Any) -> None:
 
 
 def wrap_pytest(sci: Slipcover, file_matcher: FileMatcher):
-    from . import bytecode as bc
+    def redirect_calls(module, funcName, funcWrapperName):
+        """Redirects calls to the given function to a wrapper function in the same module."""
+        import ast
+        import types
+
+        assert funcWrapperName not in module.__dict__, f"function {funcWrapperName} name already defined"
+
+        with open(module.__file__) as f:
+            t = ast.parse(f.read())
+
+        funcNames = set() # names of the functions we modified
+        for n in ast.walk(t):
+            if isinstance(n, (ast.FunctionDef, ast.AsyncFunctionDef)):
+                for s in ast.walk(n):
+                    if isinstance(s, ast.Call) and isinstance(s.func, ast.Name) and s.func.id == funcName:
+                        s.func.id = funcWrapperName
+                        funcNames.add(n.name)
+
+        code = compile(t, module.__file__, "exec")
+
+        # It's tempting to just exec(code, module.__dict__) here, but the code often times has side effects...
+        # So instead of we find the new code object(s) and replace them in the loaded module.
+
+        replacement = dict() # replacement code objects
+        def find_replacements(co):
+            for c in co.co_consts:
+                if isinstance(c, types.CodeType):
+                    if c.co_name in funcNames:
+                        replacement[c.co_name] = c
+                    find_replacements(c)
+
+        find_replacements(code)
+
+        visited = set()
+        for f in Slipcover.find_functions(module.__dict__.values(), visited):
+            if (repl := replacement.get(f.__code__.co_name, None)):
+                assert f.__code__.co_firstlineno == repl.co_firstlineno # sanity check
+                f.__code__ = repl
 
-    def exec_wrapper(obj, g):
-        if hasattr(obj, 'co_filename') and file_matcher.matches(obj.co_filename):
-            obj = sci.instrument(obj)
-        exec(obj, g)
 
     try:
         import _pytest.assertion.rewrite as pyrewrite
     except ModuleNotFoundError:
         return
 
-    for f in Slipcover.find_functions(pyrewrite.__dict__.values(), set()):
-        if 'exec' in f.__code__.co_names:
-            ed = bc.Editor(f.__code__)
-            wrapper_index = ed.add_const(exec_wrapper)
-            ed.replace_global_with_const('exec', wrapper_index)
-            f.__code__ = ed.finish()
+    redirect_calls(pyrewrite, "exec", "_Slipcover_exec_wrapper")
+
+    def exec_wrapper(obj, g):
+        if hasattr(obj, 'co_filename') and file_matcher.matches(obj.co_filename):
+            obj = sci.instrument(obj)
+        exec(obj, g)
+
+    pyrewrite._Slipcover_exec_wrapper = exec_wrapper
 
     if sci.branch:
         import inspect