| import os |
| import io |
| from contextlib import redirect_stdout |
| import pickle |
| import regex |
| import copy |
| from typing import Any, Dict, Optional |
| import multiprocess |
| from pebble import ProcessPool |
| from concurrent.futures import TimeoutError |
| from functools import partial |
| import traceback |
| from timeout_decorator import timeout |
|
|
|
|
| class GenericRuntime: |
| GLOBAL_DICT = {} |
| LOCAL_DICT = None |
| HEADERS = [] |
|
|
| def __init__(self): |
| self._global_vars = copy.copy(self.GLOBAL_DICT) |
| self._local_vars = copy.copy(self.LOCAL_DICT) if self.LOCAL_DICT else None |
|
|
| for c in self.HEADERS: |
| self.exec_code(c) |
|
|
| def exec_code(self, code_piece: str) -> None: |
| if regex.search(r"(\s|^)?input\(", code_piece) or regex.search( |
| r"(\s|^)?os.system\(", code_piece |
| ): |
| raise RuntimeError() |
| exec(code_piece, self._global_vars) |
|
|
| def eval_code(self, expr: str) -> Any: |
| return eval(expr, self._global_vars) |
|
|
| def inject(self, var_dict: Dict[str, Any]) -> None: |
| for k, v in var_dict.items(): |
| self._global_vars[k] = v |
|
|
| @property |
| def answer(self): |
| return self._global_vars["answer"] |
|
|
|
|
| class PythonExecutor: |
| def __init__( |
| self, |
| runtime: Optional[Any] = None, |
| get_answer_symbol: Optional[str] = None, |
| get_answer_expr: Optional[str] = None, |
| get_answer_from_stdout: bool = False, |
| ) -> None: |
| self.runtime = runtime if runtime else GenericRuntime() |
| self.answer_symbol = get_answer_symbol |
| self.answer_expr = get_answer_expr |
| self.get_answer_from_stdout = get_answer_from_stdout |
|
|
| def process_generation_to_code(self, gens: str): |
| batch_code = [] |
| for g in gens: |
| multiline_comments = False |
| code = [] |
| for line in g.split("\n"): |
| strip_line = line.strip() |
| if strip_line.startswith("#"): |
| line = line.split("#", 1)[0] + "# comments" |
| elif ( |
| not multiline_comments |
| and strip_line.startswith('"""') |
| and strip_line.endswith('"""') |
| and len(strip_line) >= 6 |
| ): |
| line = line.split('"""', 1)[0] + '"""comments"""' |
| elif not multiline_comments and strip_line.startswith('"""'): |
| multiline_comments = True |
| elif multiline_comments and strip_line.endswith('"""'): |
| multiline_comments = False |
| line = "" |
| if not multiline_comments: |
| code.append(line) |
| batch_code.append(code) |
| return batch_code |
|
|
| @staticmethod |
| def execute( |
| code, |
| get_answer_from_stdout=None, |
| runtime=None, |
| answer_symbol=None, |
| answer_expr=None, |
| timeout_length=10, |
| ): |
| try: |
| if get_answer_from_stdout: |
| program_io = io.StringIO() |
| with redirect_stdout(program_io): |
| timeout(timeout_length)(runtime.exec_code)("\n".join(code)) |
| program_io.seek(0) |
| result = "".join(program_io.readlines()) |
| elif answer_symbol: |
| timeout(timeout_length)(runtime.exec_code)("\n".join(code)) |
| result = runtime._global_vars[answer_symbol] |
| elif answer_expr: |
| timeout(timeout_length)(runtime.exec_code)("\n".join(code)) |
| result = timeout(timeout_length)(runtime.eval_code)(answer_expr) |
| else: |
| timeout(timeout_length)(runtime.exec_code)("\n".join(code[:-1])) |
| result = timeout(timeout_length)(runtime.eval_code)(code[-1]) |
| concise_exec_info = "" |
| exec_info = "" |
| str(result) |
| pickle.dumps(result) |
| except: |
| |
| result = "" |
| concise_exec_info = traceback.format_exc().split("\n")[-2] |
| exec_info = traceback.format_exc() |
| if ( |
| get_answer_from_stdout |
| and "exec(code_piece, self._global_vars)" in exec_info |
| ): |
| exec_info = exec_info.split("exec(code_piece, self._global_vars)")[ |
| -1 |
| ].strip() |
| msg = [] |
| for line in exec_info.split("\n"): |
| patt = regex.search( |
| r'(?P<start>.*)File "(?P<file>.*)", line (?P<lno>\d+), (?P<end>.*)', |
| line, |
| ) |
| if patt is not None: |
| if "<module>" in patt.group("end"): |
| continue |
| fname = patt.group("file") |
| if "site-packages" in fname: |
| fname = f"site-packages{fname.split('site-packages', 1)[1]}" |
| line = f'{patt.group("start")}File "{fname}", {patt.group("end")}' |
| else: |
| line = f'{patt.group("start")}{patt.group("end")}' |
| else: |
| patt = regex.search( |
| r"(?P<start>.*)(?P<file>/.*site-packages/.*\.py)(?P<end>.*)", |
| line, |
| ) |
| if patt is not None: |
| line = f'{patt.group("start")}site-packages{patt.group("file").split("site-packages", 1)[1]}{patt.group("end")}' |
| msg.append(line) |
| exec_info = "\n".join(msg) |
| return result, concise_exec_info, exec_info |
|
|
| def apply(self, code): |
| return self.batch_apply([code])[0] |
|
|
| def batch_apply(self, batch_code): |
| all_code_snippets = self.process_generation_to_code(batch_code) |
| all_exec_results = [] |
| executor = partial( |
| self.execute, |
| get_answer_from_stdout=self.get_answer_from_stdout, |
| runtime=self.runtime, |
| answer_symbol=self.answer_symbol, |
| answer_expr=self.answer_expr, |
| timeout_length=10, |
| ) |
| with ProcessPool(max_workers=multiprocess.cpu_count()) as pool: |
| iterator = pool.map(executor, all_code_snippets, timeout=10).result() |
|
|
| while True: |
| try: |
| result = next(iterator) |
| all_exec_results.append(result) |
| except StopIteration: |
| break |
| except TimeoutError as error: |
| all_exec_results.append(("", "Timeout Error", "Timeout Error")) |
| except Exception as error: |
| print(error) |
| exit() |
|
|
| batch_results = [] |
| for code, (result, concise_exec_info, exec_info) in zip( |
| all_code_snippets, all_exec_results |
| ): |
| metadata = { |
| "code": code, |
| "exec_result": result, |
| "concise_exec_info": concise_exec_info, |
| "exec_info": exec_info, |
| } |
| batch_results.append((result, metadata)) |
| return batch_results |
|
|
