- factored out code to find lines and branches in code to

  simplify upcoming changes;

src/slipcover/slipcover.py

@@ -2,7 +2,7 @@
 import sys
 import dis
 import types
-from typing import Dict, Set, List
+from typing import Dict, Set, List, Tuple
 from collections import defaultdict, Counter
 import threading
 
@@ -24,6 +24,10 @@ def counter_total(self: Counter) -> int:
     setattr(Counter, 'total', counter_total)
 
 
+if sys.version_info[0:2] >= (3,12):
+    _op_RESUME = dis.opmap["RESUME"]
+
+
 class PathSimplifier:
     def __init__(self):
         self.cwd = Path.cwd()
@@ -35,6 +39,7 @@ def simplify(self, path : str) -> str:
         except ValueError:
             return path 
 
+
 class Slipcover:
     def __init__(self, immediate: bool = False,
                  d_miss_threshold: int = 50, branch: bool = False, skip_covered: bool = False,
@@ -68,6 +73,7 @@ def handle_line(code, line):
 
             if sys.monitoring.get_tool(sys.monitoring.COVERAGE_ID) != "SlipCover":
                 sys.monitoring.use_tool_id(sys.monitoring.COVERAGE_ID, "SlipCover") # FIXME add free_tool_id
+
             sys.monitoring.register_callback(sys.monitoring.COVERAGE_ID,
                                              sys.monitoring.events.LINE, handle_line)
         else:
@@ -95,6 +101,48 @@ def _get_newly_seen(self):
         return newly_seen
 
 
+    if sys.version_info[0:2] >= (3,12):
+        @staticmethod
+        def lines_from_code(co: types.CodeType) -> Iterator[int]:
+            for c in co.co_consts:
+                if isinstance(c, types.CodeType):
+                    yield from Slipcover.lines_from_code(c)
+
+            # Python 3.11+ generates a line just for RESUME
+            yield from (line for off, line in dis.findlinestarts(co)
+                        if not br.is_branch(line) and co.co_code[off] != _op_RESUME)
+
+
+        @staticmethod
+        def branches_from_code(co: types.CodeType) -> Iterator[Tuple[int, int]]:
+            for c in co.co_consts:
+                if isinstance(c, types.CodeType):
+                    yield from Slipcover.branches_from_code(c)
+
+            yield from (br.decode_branch(line) for off, line in dis.findlinestarts(co) if br.is_branch(line))
+
+    else:
+        @staticmethod
+        def lines_from_code(co: types.CodeType) -> Iterator[int]:
+            for c in co.co_consts:
+                if isinstance(c, types.CodeType):
+                    yield from Slipcover.lines_from_code(c)
+
+            # Python 3.11 generates a 0th line; 3.11+ generates a line just for RESUME
+            yield from (line for off, line in dis.findlinestarts(co) if line != 0 and co.co_code[off] != bc.op_RESUME)
+
+
+        @staticmethod
+        def branches_from_code(co: types.CodeType) -> Iterator[Tuple[int, int]]:
+            for c in co.co_consts:
+                if isinstance(c, types.CodeType):
+                    yield from Slipcover.branches_from_code(c)
+
+            ed = bc.Editor(co)
+            for _, _, br_index in ed.find_const_assignments(br.BRANCH_NAME):
+                yield co.co_consts[br_index]
+
+
     if sys.version_info[0:2] >= (3,12):
         def instrument(self, co: types.CodeType, parent: types.CodeType = 0) -> types.CodeType:
             """Instruments a code object for coverage detection.
@@ -115,16 +163,11 @@ def instrument(self, co: types.CodeType, parent: types.CodeType = 0) -> types.Co
                 if isinstance(c, types.CodeType):
                     self.instrument(c, co)
 
-            op_RESUME = dis.opmap["RESUME"]
+            if not parent:
+                with self.lock:
+                    self.code_lines[co.co_filename].update(Slipcover.lines_from_code(co))
+                    self.code_branches[co.co_filename].update(Slipcover.branches_from_code(co))
 
-            with self.lock:
-                # Python 3.11 generates a 0th line; 3.11+ generates a line just for RESUME
-                self.code_lines[co.co_filename].update(line for off, line in dis.findlinestarts(co) \
-                                                       if line != 0 and not br.is_branch(line) \
-                                                          and co.co_code[off] != op_RESUME)
-
-                self.code_branches[co.co_filename].update(br.decode_branch(line) for off, line in dis.findlinestarts(co) \
-                                                          if br.is_branch(line) and co.co_code[off] != op_RESUME)
             return co
 
     else:
@@ -158,7 +201,6 @@ def instrument(self, co: types.CodeType, parent: types.CodeType = 0) -> types.Co
                 # are two being inserted at the same offset, the accumulated offset 'delta' applies
                 off_list.sort(key = lambda x: (x[0], len(x)))
 
-            branch_set = set()
             insert_labels = []
             probes = []
 
@@ -186,7 +228,6 @@ def instrument(self, co: types.CodeType, parent: types.CodeType = 0) -> types.Co
                     begin_off, end_off, branch_index = off_item
                     branch = co.co_consts[branch_index]
 
-                    branch_set.add(branch)
                     insert_labels.append(branch)
 
                     tr = probe.new(self, co.co_filename, branch, self.d_miss_threshold)
@@ -209,12 +250,10 @@ def instrument(self, co: types.CodeType, parent: types.CodeType = 0) -> types.Co
                 index = list(zip(ed.get_inserts(), insert_labels))
 
             with self.lock:
-                # Python 3.11 generates a 0th line; 3.11+ generates a line just for RESUME
-                self.code_lines[co.co_filename].update(line for off, line in dis.findlinestarts(co) \
-                                                       if line != 0 and co.co_code[off] != bc.op_RESUME)
-                self.code_branches[co.co_filename].update(branch_set)
-
                 if not parent:
+                    self.code_lines[co.co_filename].update(Slipcover.lines_from_code(co))
+                    self.code_branches[co.co_filename].update(Slipcover.branches_from_code(co))
+
                     self.instrumented[co.co_filename].add(new_code)
 
                 if not self.immediate:

tests/test_branch.py

@@ -12,30 +12,8 @@ def ast_parse(s):
 
 
 def get_branches(code):
-    """Extracts a list of all branches marked up in bytecode."""
-    import types
-    import dis
-
-    branches = []
-
-    # handle functions-within-functions
-    for c in code.co_consts:
-        if isinstance(c, types.CodeType):
-            branches.extend(get_branches(c))
-
-    const = None
-
-    for inst in dis.get_instructions(code):
-        if inst.opname == 'LOAD_CONST':
-            const = inst.argval
-
-        elif const is not None:
-            if inst.opname in ('STORE_NAME', 'STORE_GLOBAL') and inst.argval == br.BRANCH_NAME:
-                branches.append(const)
-
-            const = None
-
-    return sorted(branches)
+    from slipcover.slipcover import Slipcover
+    return sorted(Slipcover.branches_from_code(code))
 
 
 def assign2append(tree: ast.AST):