fix: execute js/py code (#129)

2 files changed, 28 insertions, 8 deletions

diff --git a/tools/execute_js_code.js b/tools/execute_js_code.js
index 6f04ade..0b7557b 100644
--- a/tools/execute_js_code.js
+++ b/tools/execute_js_code.js
@@ -12,8 +12,11 @@ exports.run = function run({ code }) {
     if (callback) callback();
   };
 
-  eval(code);
-  
+  const value = eval(code);
+  if (value !== undefined) {
+    output += value;
+  }
+
   process.stdout.write = oldStdoutWrite;
   return output;
 }
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)