fix: execute js/py code (#129)

1 files changed, 23 insertions, 6 deletions

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)