Auto-run zero-arg Python functions in sandbox

app.py

@@ -1,6 +1,7 @@
 import json
 import os
 import re
+import ast
 from pathlib import Path
 
 import gradio as gr
@@ -266,7 +267,7 @@ async def run_sandbox(language: str, draft_code: str, verdict_json: str) -> str:
     if language.lower() != "python":
         return "Sandbox execution is currently wired for Python only."
 
-    code = code_from_verdict(draft_code, verdict_json)
+    code = prepare_python_for_sandbox(code_from_verdict(draft_code, verdict_json))
     if not code.strip():
         return "No generated code is available yet."
 
@@ -287,6 +288,52 @@ async def run_sandbox(language: str, draft_code: str, verdict_json: str) -> str:
     )
 
 
+def prepare_python_for_sandbox(code: str) -> str:
+    code = strip_markdown_code_fence(code)
+    try:
+        tree = ast.parse(code)
+    except SyntaxError:
+        return code
+
+    executable_nodes = (
+        ast.Assign,
+        ast.AugAssign,
+        ast.AnnAssign,
+        ast.Assert,
+        ast.Delete,
+        ast.Expr,
+        ast.For,
+        ast.AsyncFor,
+        ast.While,
+        ast.If,
+        ast.Match,
+        ast.Raise,
+        ast.Return,
+        ast.Try,
+        ast.With,
+        ast.AsyncWith,
+    )
+    has_top_level_execution = any(isinstance(node, executable_nodes) for node in tree.body)
+    if has_top_level_execution:
+        return code
+
+    for node in tree.body:
+        if isinstance(node, ast.FunctionDef) and function_has_no_required_args(node):
+            return f'{code}\n\nif __name__ == "__main__":\n    {node.name}()\n'
+
+    return code
+
+
+def function_has_no_required_args(node: ast.FunctionDef) -> bool:
+    args = node.args
+    positional = [*args.posonlyargs, *args.args]
+    required_positional = len(positional) - len(args.defaults)
+    required_kwonly = sum(
+        1 for arg, default in zip(args.kwonlyargs, args.kw_defaults) if default is None
+    )
+    return required_positional == 0 and required_kwonly == 0
+
+
 with gr.Blocks(
     title="Split-Brain Co-Pilot",
     css=load_static("style.css"),