| import os |
| import pickle |
| import sys |
| from importlib import import_module |
| from io import StringIO |
| from typing import Any, Dict, List |
|
|
| import coverage |
| from rich.progress import track |
|
|
| from evalplus.eval.utils import swallow_io |
| from tools.tsr.utils import get_problems, get_task_ids, to_path |
|
|
|
|
| class Capturing(list): |
| def __enter__(self): |
| self._stdout = sys.stdout |
| sys.stdout = self._stringio = StringIO() |
| return self |
|
|
| def __exit__(self, *args): |
| self.extend(self._stringio.getvalue().splitlines()) |
| del self._stringio |
| sys.stdout = self._stdout |
|
|
|
|
| def parse_lcov(outputs: List[str]): |
| switch, extracted_outputs = False, [] |
| for line in outputs: |
| if switch == False and "tmp_src" in line: |
| switch = True |
| if switch == True and "end_of_record" in line: |
| switch = False |
| if switch: |
| extracted_outputs.append(line) |
|
|
| branch, branch_covered = [], [] |
| for line in extracted_outputs: |
| if line.startswith("BRDA"): |
| |
| lineno, blockno, branchno, taken = line[5:].split(",") |
| branch_sig = f"BR:{lineno},{blockno},{branchno}" |
| branch.append(branch_sig) |
| if taken not in ["0", "-"]: |
| branch_covered.append(branch_sig) |
| per = 1.0 if len(branch) == 0 else len(branch_covered) / len(branch) |
| return per, branch, branch_covered |
|
|
|
|
| def test_code_coverage( |
| identifier: str, code: str, inputs: List[List[Any]], entry_point: str |
| ): |
| module_name = f"tmp_src_{identifier}" |
| with open(f"{module_name}.py", "w") as f: |
| f.write(code) |
|
|
| mod = import_module(module_name) |
| func = getattr(mod, entry_point, None) |
| assert func != None, f"entry_point = {entry_point} not exist, code: {code}" |
|
|
| cov = coverage.Coverage(branch=True) |
| cov.start() |
| with swallow_io(): |
| for input_list in inputs: |
| func(*input_list) |
| cov.stop() |
| with Capturing() as outputs: |
| cov.lcov_report(outfile="-") |
|
|
| ret = parse_lcov(outputs) |
|
|
| os.remove(f"{module_name}.py") |
| return ret |
|
|
|
|
| def collect_coverage_info(coverage_dir: str, dataset: str) -> Dict[str, Dict[str, Any]]: |
| os.makedirs(coverage_dir, exist_ok=True) |
| problems = get_problems(dataset) |
| task_ids = get_task_ids(dataset) |
| coverage_info = {task_id: {} for task_id in task_ids} |
| for task_id in track(task_ids, description="Testing gt coverage..."): |
| coverage_cache_path = os.path.join(coverage_dir, f"{to_path(task_id)}.pkl") |
| if os.path.isfile(coverage_cache_path): |
| with open(coverage_cache_path, "rb") as f: |
| coverage_info[task_id] = pickle.load(f) |
| continue |
| groundtruth_code = ( |
| problems[task_id]["prompt"] + problems[task_id]["canonical_solution"] |
| ) |
| plus_tests = problems[task_id]["plus_input"] |
| entry_point = problems[task_id]["entry_point"] |
| for i, plus_test in enumerate(plus_tests): |
| per, branch, branch_covered = test_code_coverage( |
| to_path(task_id), groundtruth_code, [plus_test], entry_point |
| ) |
| test_id = f"plus_{i}" |
| coverage_info[task_id].setdefault(test_id, []).extend( |
| [(br, "gt") for br in branch_covered] |
| ) |
| with open(coverage_cache_path, "wb") as f: |
| pickle.dump(coverage_info[task_id], f) |
|
|
| return coverage_info |
|
|
|
|
| if __name__ == "__main__": |
| import argparse |
|
|
| parser = argparse.ArgumentParser() |
| parser.add_argument("--dataset", type=str, choices=["humaneval", "mbpp"]) |
| parser.add_argument("--report_dir", required=True, type=str) |
| args = parser.parse_args() |
|
|
| coverage_dir = os.path.join(args.report_dir, "coverage_cache") |
| collect_coverage_info(coverage_dir, args.dataset) |
|
|
