From ecf7233401ebe273fa292b8208651395d99ddff9 Mon Sep 17 00:00:00 2001 From: sigoden Date: Mon, 25 Nov 2024 08:21:58 +0800 Subject: fix: execute js/py code (#129) --- tools/execute_py_code.py | 29 +++++++++++++++++++++++------ 1 file changed, 23 insertions(+), 6 deletions(-) (limited to 'tools/execute_py_code.py') diff --git a/tools/execute_py_code.py b/tools/execute_py_code.py index a8cfe48..71d66e1 100644 --- a/tools/execute_py_code.py +++ b/tools/execute_py_code.py @@ -1,16 +1,33 @@ +import ast import io -import sys +from contextlib import redirect_stdout + def run(code: str): """Execute the python code. Args: code: Python code to execute, such as `print("hello world")` """ - old_stdout = sys.stdout output = io.StringIO() - sys.stdout = output + with redirect_stdout(output): + value = exec_with_return(code, {}, {}) + + if value is not None: + output.write(str(value)) + + return output.getvalue() - exec(code) - sys.stdout = old_stdout - return output.getvalue() \ No newline at end of file +def exec_with_return(code: str, globals: dict, locals: dict): + a = ast.parse(code) + last_expression = None + if a.body: + if isinstance(a_last := a.body[-1], ast.Expr): + last_expression = ast.unparse(a.body.pop()) + elif isinstance(a_last, ast.Assign): + last_expression = ast.unparse(a_last.targets[0]) + elif isinstance(a_last, (ast.AnnAssign, ast.AugAssign)): + last_expression = ast.unparse(a_last.target) + exec(ast.unparse(a), globals, locals) + if last_expression: + return eval(last_expression, globals, locals) -- cgit v1.2.3