| 50 | |
| 51 | |
| 52 | def get_code_and_contract_and_assertion(task_id: id) -> Tuple[str, str, str]: |
| 53 | py_file_path = str(GROUNDTRUTH_MBPP_PATH) + f"/{str(task_id).zfill(3)}.py" |
| 54 | with open(py_file_path) as reader: |
| 55 | text = reader.read() |
| 56 | # remove docstring |
| 57 | start_index = text.find('"""') |
| 58 | end_index = text.find('"""', start_index + 3) |
| 59 | if start_index != -1 and end_index != -1: |
| 60 | text = text[:start_index] + text[end_index + 3 :] |
| 61 | |
| 62 | lines = text.splitlines() |
| 63 | assertion = "" |
| 64 | contract = "" |
| 65 | |
| 66 | for i in range(len(lines)): |
| 67 | if "$_CONTRACT_$" in lines[i]: |
| 68 | contract += lines[i] + "\n" |
| 69 | elif lines[i].startswith("assert"): |
| 70 | assertion += lines[i] + "\n" |
| 71 | |
| 72 | for i in range(len(lines) - 1, -1, -1): |
| 73 | if ( |
| 74 | "$_CONTRACT_$" in lines[i] |
| 75 | or lines[i].startswith("assert") |
| 76 | or lines[i] == "" |
| 77 | ): |
| 78 | del lines[i] |
| 79 | |
| 80 | for i in range(len(lines) - 1, -1, -1): |
| 81 | if lines[i].startswith("import"): |
| 82 | del lines[i] |
| 83 | else: |
| 84 | break |
| 85 | |
| 86 | code = "\n".join(lines) |
| 87 | return "\n" + code + "\n", "\n" + contract, "\n" + assertion |
| 88 | |
| 89 | |
| 90 | def instrument_inputs(code, entry_point, test_code) -> str: |