File size: 7,253 Bytes
22c93a7 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 |
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()) # [-1]
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) # serialization check
except:
# traceback.print_exc()
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
|