mirror of
https://github.com/EvolutionAPI/adk-python.git
synced 2025-07-16 04:02:55 -06:00
fix: add support for running python main function in UnsafeLocalCodeExecutor when the code has an if __name__ == "__main__" statement.
Fixes https://github.com/google/adk-python/issues/791 PiperOrigin-RevId: 765359512
This commit is contained in:
parent
4075290a1d
commit
95e33baf57
@ -12,8 +12,12 @@
|
|||||||
# See the License for the specific language governing permissions and
|
# See the License for the specific language governing permissions and
|
||||||
# limitations under the License.
|
# limitations under the License.
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
from contextlib import redirect_stdout
|
from contextlib import redirect_stdout
|
||||||
import io
|
import io
|
||||||
|
import re
|
||||||
|
from typing import Any
|
||||||
|
|
||||||
from pydantic import Field
|
from pydantic import Field
|
||||||
from typing_extensions import override
|
from typing_extensions import override
|
||||||
@ -24,6 +28,12 @@ from .code_execution_utils import CodeExecutionInput
|
|||||||
from .code_execution_utils import CodeExecutionResult
|
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):
|
class UnsafeLocalCodeExecutor(BaseCodeExecutor):
|
||||||
"""A code executor that unsafely execute code in the current local context."""
|
"""A code executor that unsafely execute code in the current local context."""
|
||||||
|
|
||||||
@ -55,6 +65,7 @@ class UnsafeLocalCodeExecutor(BaseCodeExecutor):
|
|||||||
error = ''
|
error = ''
|
||||||
try:
|
try:
|
||||||
globals_ = {}
|
globals_ = {}
|
||||||
|
_prepare_globals(code_execution_input.code, globals_)
|
||||||
locals_ = {}
|
locals_ = {}
|
||||||
stdout = io.StringIO()
|
stdout = io.StringIO()
|
||||||
with redirect_stdout(stdout):
|
with redirect_stdout(stdout):
|
||||||
|
Loading…
Reference in New Issue
Block a user