From 44018e369038c8aa0dfc4a9ffc52d8735f074347 Mon Sep 17 00:00:00 2001 From: Juan Altmayer Pizzorno Date: Thu, 26 Oct 2023 19:44:38 -0400 Subject: [PATCH] - factored out code to find lines and branches in code to simplify upcoming changes; --- src/slipcover/slipcover.py | 73 +++++++++++++++++++++++++++++--------- tests/test_branch.py | 26 ++------------ 2 files changed, 58 insertions(+), 41 deletions(-) diff --git a/src/slipcover/slipcover.py b/src/slipcover/slipcover.py index 26e8512..9791d88 100644 --- a/src/slipcover/slipcover.py +++ b/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: diff --git a/tests/test_branch.py b/tests/test_branch.py index 15a25b1..2024044 100644 --- a/tests/test_branch.py +++ b/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):