mirror of
https://github.com/agentscope-ai/ReMe.git
synced 2026-08-28 05:25:04 +00:00
142 lines
4.3 KiB
Python
142 lines
4.3 KiB
Python
"""Utility functions for executing code and shell commands.
|
|
|
|
This module provides helper functions for running Python code and shell commands,
|
|
with support for async execution and output capture.
|
|
"""
|
|
|
|
import asyncio
|
|
import concurrent.futures
|
|
import contextlib
|
|
from io import StringIO
|
|
|
|
|
|
async def run_shell_command(cmd: str, timeout: float | None = 30) -> tuple[str, str, int]:
|
|
"""Execute a shell command asynchronously.
|
|
|
|
Args:
|
|
cmd: The shell command to execute.
|
|
timeout: Maximum time to wait for command completion in seconds. None for no timeout.
|
|
|
|
Returns:
|
|
A tuple containing (stdout, stderr, return_code) as strings and integer.
|
|
|
|
Raises:
|
|
TimeoutError: If the command does not complete within the timeout.
|
|
"""
|
|
process = await asyncio.create_subprocess_shell(
|
|
cmd,
|
|
stdout=asyncio.subprocess.PIPE,
|
|
stderr=asyncio.subprocess.PIPE,
|
|
)
|
|
|
|
try:
|
|
if timeout:
|
|
stdout, stderr = await asyncio.wait_for(process.communicate(), timeout=timeout)
|
|
else:
|
|
stdout, stderr = await process.communicate()
|
|
except (asyncio.TimeoutError, TimeoutError) as e:
|
|
# Kill the child process to avoid orphaned / zombie processes
|
|
process.kill()
|
|
await process.wait()
|
|
raise TimeoutError(f"Shell command timed out after {timeout}s") from e
|
|
|
|
return (
|
|
stdout.decode("utf-8", errors="ignore"),
|
|
stderr.decode("utf-8", errors="ignore"),
|
|
process.returncode,
|
|
)
|
|
|
|
|
|
def exec_code(
|
|
code: str,
|
|
timeout: float | None = 30,
|
|
executor: concurrent.futures.ThreadPoolExecutor | None = None,
|
|
) -> str:
|
|
"""Execute Python code and capture the output.
|
|
|
|
Args:
|
|
code: The Python code string to execute.
|
|
timeout: Maximum time to wait for execution in seconds. None for no timeout.
|
|
executor: Optional thread pool executor to use. If None, a temporary
|
|
single-thread executor is created (and shut down after the call).
|
|
Pass a shared executor to amortize thread-creation overhead across
|
|
multiple calls.
|
|
|
|
Returns:
|
|
The captured stdout output, or the error message if execution fails.
|
|
|
|
Raises:
|
|
TimeoutError: If execution exceeds the timeout.
|
|
"""
|
|
|
|
def _run() -> str:
|
|
redirected_output = StringIO()
|
|
with contextlib.redirect_stdout(redirected_output):
|
|
exec(code)
|
|
return redirected_output.getvalue()
|
|
|
|
def _submit(pool: concurrent.futures.ThreadPoolExecutor) -> str:
|
|
future = pool.submit(_run)
|
|
return future.result(timeout=timeout)
|
|
|
|
try:
|
|
if executor is not None:
|
|
return _submit(executor)
|
|
with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool:
|
|
return _submit(pool)
|
|
|
|
except concurrent.futures.TimeoutError as e:
|
|
raise TimeoutError(f"Code execution timed out after {timeout}s") from e
|
|
|
|
except Exception as e:
|
|
return str(e)
|
|
|
|
except BaseException as e:
|
|
return str(e)
|
|
|
|
|
|
async def async_exec_code(
|
|
code: str,
|
|
timeout: float | None = 30,
|
|
executor: concurrent.futures.ThreadPoolExecutor | None = None,
|
|
) -> str:
|
|
"""Execute Python code asynchronously and capture the output.
|
|
|
|
Runs the code in a thread executor to avoid blocking the event loop,
|
|
with async-friendly timeout via ``asyncio.wait_for``.
|
|
|
|
Args:
|
|
code: The Python code string to execute.
|
|
timeout: Maximum time to wait for execution in seconds. None for no timeout.
|
|
executor: Optional thread pool executor. If None, the default event-loop
|
|
executor is used.
|
|
|
|
Returns:
|
|
The captured stdout output, or the error message if execution fails.
|
|
|
|
Raises:
|
|
TimeoutError: If execution exceeds the timeout.
|
|
"""
|
|
|
|
def _run() -> str:
|
|
redirected_output = StringIO()
|
|
with contextlib.redirect_stdout(redirected_output):
|
|
exec(code)
|
|
return redirected_output.getvalue()
|
|
|
|
loop = asyncio.get_running_loop()
|
|
|
|
try:
|
|
coro = loop.run_in_executor(executor, _run)
|
|
if timeout is not None:
|
|
return await asyncio.wait_for(coro, timeout=timeout)
|
|
return await coro
|
|
|
|
except asyncio.TimeoutError as e:
|
|
raise TimeoutError(f"Code execution timed out after {timeout}s") from e
|
|
|
|
except Exception as e:
|
|
return str(e)
|
|
|
|
except BaseException as e:
|
|
return str(e)
|