refactor: nicer branch_trails code

Author: nedbat

coverage/bytecode.py

@@ -63,7 +63,6 @@ class InstructionWalker:
     First, in strict sequence to visit all the instructions in the object.
     This is `walk(follow_jumps=False)`.  Second, we want to follow jumps to
     understand how execution will flow: `walk(follow_jumps=True)`.
-
     """
 
     def __init__(self, code: CodeType) -> None:
@@ -98,7 +97,8 @@ def walk(
             offset += 2
 
 
-TBranchTrails = dict[TOffset, dict[Optional[TArc], set[TOffset]]]
+TBranchTrailsOneSource = dict[Optional[TArc], set[TOffset]]
+TBranchTrails = dict[TOffset, TBranchTrailsOneSource]
 
 
 def branch_trails(code: CodeType) -> TBranchTrails:
@@ -117,7 +117,9 @@ def branch_trails(code: CodeType) -> TBranchTrails:
     arc from the original instruction's line to the new source line.
 
     """
-    the_trails: TBranchTrails = collections.defaultdict(lambda:collections.defaultdict(set))
+    the_trails: TBranchTrails = collections.defaultdict(
+        lambda: collections.defaultdict(set)
+    )
     iwalker = InstructionWalker(code)
     for inst in iwalker.walk(follow_jumps=False):
         if not inst.jump_target:
@@ -131,7 +133,10 @@ def branch_trails(code: CodeType) -> TBranchTrails:
         if from_line is None:
             continue
 
-        def walk_one_branch(start_at: TOffset) -> tuple[Optional[TArc], set[TOffset]]:
+        def add_one_branch_trail(
+            trails: TBranchTrailsOneSource,
+            start_at: TOffset,
+        ) -> None:
             # pylint: disable=cell-var-from-loop
             inst_offsets: set[TOffset] = set()
             to_line = None
@@ -146,17 +151,15 @@ def walk_one_branch(start_at: TOffset) -> tuple[Optional[TArc], set[TOffset]]:
                     to_line = -code.co_firstlineno
                     break
             if to_line is not None:
-                return (from_line, to_line), inst_offsets
+                trails[(from_line, to_line)].update(inst_offsets)
             else:
-                return None, set()
+                trails[None] = set()
 
         # Calculate two trails: one from the next instruction, and one from the
         # jump_target instruction.
-        trails = collections.defaultdict(set)
-        arc, offsets = walk_one_branch(start_at=inst.offset + 2)
-        trails[arc].update(offsets)
-        arc, offsets = walk_one_branch(start_at=inst.jump_target)
-        trails[arc].update(offsets)
+        trails: TBranchTrailsOneSource = collections.defaultdict(set)
+        add_one_branch_trail(trails, start_at=inst.offset + 2)
+        add_one_branch_trail(trails, start_at=inst.jump_target)
         the_trails[inst.offset] = trails
 
         # Sometimes we get BRANCH_RIGHT or BRANCH_LEFT events from instructions

tests/helpers.py

@@ -396,16 +396,15 @@ def all_our_source_files() -> Iterator[tuple[Path, str]]:
 
     Produces a stream of (filename, file contents) tuples.
     """
-    print(f"all_our_source_files: {coverage.__file__ = }")
     cov_dir = Path(coverage.__file__).parent.parent
     if ".tox" in cov_dir.parts:
         # We are in a tox-installed environment, look above the .tox dir to
         # also find the uninstalled source files.
         cov_dir = Path(os.fspath(cov_dir).partition(".tox")[0])
 
-    print(f"all_our_source_files: {os.path.abspath(cov_dir) = }")
     # To run against all the files in the tox venvs:
     #   for source_file in cov_dir.rglob("*.py"):
     for sub in [".", "benchmark", "ci", "coverage", "lab", "tests"]:
+        assert (cov_dir / sub).is_dir()
         for source_file in (cov_dir / sub).glob("*.py"):
             yield (source_file, source_file.read_text(encoding="utf-8"))

tests/test_arcs.py

@@ -2129,7 +2129,10 @@ def func():
             branchz_missing="29 38 45 56 5. 9A 9.",
         )
 
-    @pytest.mark.skipif(env.PYVERSION[:2] == (3, 9), reason="not sure why it isn't fixed on 3.9")
+    @pytest.mark.skipif(
+        env.PYVERSION[:2] == (3, 9),
+        reason="CPython fix not backported to 3.9: https://github.com/python/cpython/issues/93061",
+    )
     def test_bug_1999(self) -> None:
         self.check_coverage("""\
             import asyncio