fix: add support for running python main function in UnsafeLocalCodeExecutor when the code has an if __name__ == "__main__" statement.

google-genai-bot · copybara-github · commit 95e33baf57e9 · 2025-05-30T15:30:01.000-07:00
Fixes google#791 PiperOrigin-RevId: 765359512
diff --git a/src/google/adk/code_executors/unsafe_local_code_executor.py b/src/google/adk/code_executors/unsafe_local_code_executor.py
@@ -12,8 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 from contextlib import redirect_stdout
 import io
+import re
+from typing import Any
 
 from pydantic import Field
 from typing_extensions import override
@@ -24,6 +28,12 @@
 from .code_execution_utils import CodeExecutionResult
 
 
+def _prepare_globals(code: str, globals_: dict[str, Any]) -> None:
+  """Prepare globals for code execution, injecting __name__ if needed."""
+  if re.search(r"if\s+__name__\s*==\s*['\"]__main__['\"]", code):
+    globals_['__name__'] = '__main__'
+
+
 class UnsafeLocalCodeExecutor(BaseCodeExecutor):
   """A code executor that unsafely execute code in the current local context."""
 
@@ -55,6 +65,7 @@ def execute_code(
     error = ''
     try:
       globals_ = {}
+      _prepare_globals(code_execution_input.code, globals_)
       locals_ = {}
       stdout = io.StringIO()
       with redirect_stdout(stdout):
