pygx.coding

Code generation utilities.

coding

Code generation utilities.

pygx.coding provides a small toolkit for compiling and executing Python source code at runtime under controlled conditions. It is used internally by PyGX to materialize generated programs, but is exposed as a public sub-package because the building blocks are useful on their own:

• parse / evaluate / run for compiling and executing Python source.
• make_function for synthesizing a callable from a source string.
• permission / get_permission for declaring which language features (assignment, function definition, import) the executed source may use.
• sandbox_call / maybe_sandbox_call for executing a callable in a child process with a timeout.
• CodeError / SerializationError for errors raised by the above APIs.

CodeError

CodeError(code: str, cause: BaseException)

Bases: RuntimeError

Python code error.

Source code in pygx/coding/_errors.py

def __init__(
    self,
    code: str,
    cause: BaseException,
):
    self.code = code
    self.cause = cause

    # Figure out the starting and ending line numbers of the erratic code.
    lineno = None
    end_lineno = None
    if isinstance(cause, SyntaxError):
        lineno = cause.lineno
        end_lineno = cause.end_lineno
    elif not isinstance(cause, TimeoutError):
        tb = sys.exc_info()[2]
        frames = traceback.extract_tb(tb, limit=5)
        for f in frames:
            if not f.filename or f.filename == '<string>':
                lineno = f.lineno
                end_lineno = lineno
                break
    self.lineno = lineno
    self.end_lineno = end_lineno

code_lines

code_lines(start_line: int, end_line: int | None)

Returns code lines .

Source code in pygx/coding/_errors.py

def code_lines(self, start_line: int, end_line: int | None):
    """Returns code lines ."""
    return '\n'.join(self.code.split('\n')[start_line:end_line])

format

format(include_complete_code: bool = True)

Formats the code error.

Source code in pygx/coding/_errors.py

def format(self, include_complete_code: bool = True):
    """Formats the code error."""
    r = io.StringIO()
    error_message = str(self.cause).rstrip()
    if 'line' not in error_message and self.lineno is not None:
        error_message += f' (<unknown>, line {self.lineno})'
    r.write(
        formatting.colored(
            f'{self.cause.__class__.__name__}: {error_message}', 'magenta'
        )
    )

    if self.lineno is not None:
        r.write('\n\n')
        r.write(
            textwrap.indent(
                formatting.colored(
                    self.code_lines(self.lineno - 1, self.end_lineno),
                    'magenta',
                ),
                ' ' * 2,
            )
        )
        r.write('\n')

    if include_complete_code:
        r.write('\n')
        r.write(formatting.colored('[Code]', 'green', styles=['bold']))
        r.write('\n\n')
        r.write(formatting.colored('  ```python\n', 'green'))
        r.write(
            textwrap.indent(formatting.colored(self.code, 'green'), ' ' * 2)
        )
        r.write(formatting.colored('\n  ```\n', 'green'))
    return r.getvalue()

SerializationError

SerializationError(message: str | None, cause: BaseException)

Bases: RuntimeError

Object serialization error.

Source code in pygx/coding/_errors.py

def __init__(self, message: str | None, cause: BaseException):
    self.message = message
    self.cause = cause

CodePermission

Bases: Flag

Permissions for code execution.

context

context(**kwargs)

Context manager to inject symbols for code execution.

Source code in pygx/coding/_execution.py

@contextlib.contextmanager
def context(**kwargs):
    """Context manager to inject symbols for code execution."""
    ctx = get_context()
    ctx.update(kwargs)
    concurrent.thread_local_push(_TLS_CODE_RUN_CONTEXT, ctx)

    try:
        yield ctx
    finally:
        concurrent.thread_local_pop(_TLS_CODE_RUN_CONTEXT)

evaluate

evaluate(
    code: str,
    *,
    global_vars: dict[str, Any] | None = None,
    permission: CodePermission | None = None,
    returns_stdout: bool = False,
    outputs_intermediate: bool = False
) -> Any | dict[str, Any]

Executes Python code.

Features

• Fine-grained execution policy for limiting what APIs could be executed. This eliminates the need for sandboxing.
• It exposes both the final results and intermediate results (variables).

Parameters:

Name	Type	Description	Default
code	str	Python code to run.	required
global_vars	dict[str, Any] | None	An optional dict as the globals that could be referenced by the code.	None
permission	CodePermission | None	Permission for the Python code to run.	None
returns_stdout	bool	If True, the stdout (a str) will be returned.	False
outputs_intermediate	bool	Applicable when returns_stdout is False. If True, intermediate output will be outputted as a dict, with the last line's value accessible by key 'result' and the std output accessible by key 'stdout'. Otherwise the value of the last line will be returned.	False

Returns:

Type	Description
Any | dict[str, Any]	The value of the last line of the code block. Or a dict of variable
Any | dict[str, Any]	names of all locals to their evaluated values as the output of the code to
Any | dict[str, Any]	run. The value for the last line can be accessed by key 'result'. Or the
Any | dict[str, Any]	stdout as a str.

Source code in pygx/coding/_execution.py

def evaluate(
    code: str,
    *,
    global_vars: dict[str, Any] | None = None,
    permission: permissions.CodePermission | None = None,
    returns_stdout: bool = False,
    outputs_intermediate: bool = False,
) -> Any | dict[str, Any]:
    """Executes Python code.

    Features:
      * Fine-grained execution policy for limiting what APIs could be executed.
        This eliminates the need for sandboxing.
      * It exposes both the final results and intermediate results (variables).

    Args:
      code: Python code to run.
      global_vars: An optional dict as the globals that could be referenced by the
        code.
      permission: Permission for the Python code to run.
      returns_stdout: If True, the stdout (a str) will be returned.
      outputs_intermediate: Applicable when returns_stdout is False. If True,
        intermediate output will be outputted as a dict, with the last line's
        value accessible by key '__result__' and the std output accessible by
        key '__stdout__'. Otherwise the value of the last line will be returned.

    Returns:
      The value of the last line of the code block. Or a dict of variable
      names of all locals to their evaluated values as the output of the code to
      run. The value for the last line can be accessed by key '__result__'. Or the
      stdout as a str.
    """
    # Set up the permission and context.
    permission = permission or permissions.get_permission()
    ctx = dict(get_context())
    if global_vars:
        ctx.update(global_vars)

    # Parse the code str.
    code_block = parsing.parse(code, permission)
    global_vars, orig_global_vars = ctx, ctx.copy()

    # No code.
    if not code_block.body:  # pytype: disable=attribute-error
        return {} if outputs_intermediate else None

    stdout = io.StringIO()
    with contextlib.redirect_stdout(stdout):
        if hasattr(
            code_block.body[-1], 'value'
        ):  # pytype: disable=attribute-error
            last_expr = code_block.body.pop()  # pytype: disable=attribute-error
            result_vars = [RESULT_KEY]

            if isinstance(last_expr, ast.Assign):
                for name_node in last_expr.targets:
                    if isinstance(name_node, ast.Name):
                        result_vars.append(name_node.id)

            last_expr = ast.Expression(getattr(last_expr, 'value'))

            try:
                # Execute the lines before the last expression.
                # NOTE(daiyip): Only a `globals` dict is specified here, which will also
                # be used to output intermediate values by `exec`. We do not specify a
                # separate `locals` dict here, for - "If exec gets two separate objects
                # as globals and locals, the code will be executed as if it were
                # embedded in a class definition." - as the Python document explains.
                # The outcome is that new functions defined in the code block could not
                # be called by other newly defined functions.
                # Refer to https://stackoverflow.com/questions/
                # 73940751/why-cant-i-call-a-function-from-another-function-using-exec
                # for more details.
                exec(compile(code_block, '', mode='exec'), global_vars)  # pylint: disable=exec-used

                # Evaluate the last expression.
                result = eval(  # pylint: disable=eval-used
                    compile(last_expr, '', mode='eval'), global_vars
                )
            except BaseException as e:
                raise errors.CodeError(code, e) from e

            for result_var in result_vars:
                global_vars[result_var] = result
        else:
            try:
                exec(compile(code_block, '', mode='exec'), global_vars)  # pylint: disable=exec-used
            except BaseException as e:
                raise errors.CodeError(code, e) from e
            global_vars[RESULT_KEY] = list(global_vars.values())[-1]

    if returns_stdout:
        return stdout.getvalue()
    if outputs_intermediate:
        outputs = {}
        for k, v in global_vars.items():
            if k == '__builtins__':
                continue
            if k not in orig_global_vars or v is not orig_global_vars[k]:
                outputs[k] = v
        # Add stdout to outputs.
        outputs[STDOUT_KEY] = stdout.getvalue()
        return outputs
    return global_vars[RESULT_KEY]

get_context

get_context() -> dict[str, Any]

Gets the current context for code execution.

Source code in pygx/coding/_execution.py

def get_context() -> dict[str, Any]:
    """Gets the current context for code execution."""
    context_stack = concurrent.thread_local_get(_TLS_CODE_RUN_CONTEXT, None)
    return dict(context_stack[-1]) if context_stack else {}

maybe_sandbox_call

maybe_sandbox_call(
    func: Callable[..., Any],
    *args: Any,
    sandbox: bool | None = None,
    timeout: float | None = None,
    **kwargs: Any
) -> Any

Maybe calls a function with sandboxing.

Parameters:

Name	Type	Description	Default
func	Callable[..., Any]	Function to call.	required
*args	Any	Postional args that will be passed to func.	()
sandbox	bool | None	If True, run code in sandbox; If False, run code in current process. If None, run in sandbox first, if the output could not be serialized and pass to current process, run the code again in current process.	None
timeout	float | None	Execution timeout in seconds. If None, wait the code the complete.	None
**kwargs	Any	Keyword args that will be passed to func.	{}

Returns:

Type	Description
Any	The return value of func.

Raises:

Type	Description
TimeoutError	If the execution time exceeds the timeout.
Exception	Exception that are raised from func.

Source code in pygx/coding/_execution.py

def maybe_sandbox_call(
    func: Callable[..., Any],
    *args: Any,
    sandbox: bool | None = None,
    timeout: float | None = None,
    **kwargs: Any,
) -> Any:
    """Maybe calls a function with sandboxing.

    Args:
      func: Function to call.
      *args: Postional args that will be passed to `func`.
      sandbox: If True, run code in sandbox; If False, run code in current
        process. If None, run in sandbox first, if the output could not be
        serialized and pass to current process, run the code again in current
        process.
      timeout: Execution timeout in seconds. If None, wait the code the complete.
      **kwargs: Keyword args that will be passed to `func`.

    Returns:
      The return value of `func`.

    Raises:
      TimeoutError: If the execution time exceeds the timeout.
      Exception: Exception  that are raised from `func`.
    """
    if sandbox is None:
        try:
            return sandbox_call(func, *args, timeout=timeout, **kwargs)
        # NOTE(daiyip): output could be serialized across processes, giving it
        # already finishes on sandbox, so it should be much safer to run under
        # current process.
        except errors.SerializationError:
            return func(*args, **kwargs)
    elif sandbox:
        return sandbox_call(func, *args, timeout=timeout, **kwargs)
    else:
        return func(*args, **kwargs)

run

run(
    code: str,
    *,
    global_vars: dict[str, Any] | None = None,
    permission: CodePermission | None = None,
    returns_stdout: bool = False,
    outputs_intermediate: bool = False,
    sandbox: bool | None = None,
    timeout: float | None = None
) -> Any | dict[str, Any]

Executes Python code.

Features

• Fine-grained execution policy for limiting what APIs could be executed. This eliminates the need for sandboxing.
• It exposes both the final results and intermediate results (variables).

Parameters:

Name	Type	Description	Default
code	str	Python code to run.	required
global_vars	dict[str, Any] | None	An optional dict of	None
permission	CodePermission | None	Permission for the Python code to run.	None
returns_stdout	bool	If True, the stdout (a str) will be returned.	False
outputs_intermediate	bool	Applicable when returns_stdout is False. If True, intermediate output will be outputted as a dict, with the last line's value accessible by key 'result' and the std output accessible by key 'stdout'. Otherwise the value of the last line will be returned.	False
sandbox	bool | None	If True, run code in sandbox; If False, run code in current process. If None, run in sandbox first, if the output could not be serialized and pass to current process, run the code again in current process.	None
timeout	float | None	Execution timeout in seconds. If None, wait the code the complete.	None

Returns:

Type	Description
Any | dict[str, Any]	The value of the last line of the code block. Or a dict of variable
Any | dict[str, Any]	names of all locals to their evaluated values as the output of the code to
Any | dict[str, Any]	run. The value for the last line can be accessed by key 'result'. Or the
Any | dict[str, Any]	stdout as a str.

Raises:

Type	Description
TimeoutError	If the execution time exceeds the timeout.
Exception	Exception that are raised from the code.

Source code in pygx/coding/_execution.py

def run(
    code: str,
    *,
    global_vars: dict[str, Any] | None = None,
    permission: permissions.CodePermission | None = None,
    returns_stdout: bool = False,
    outputs_intermediate: bool = False,
    sandbox: bool | None = None,
    timeout: float | None = None,
) -> Any | dict[str, Any]:
    """Executes Python code.

    Features:
      * Fine-grained execution policy for limiting what APIs could be executed.
        This eliminates the need for sandboxing.
      * It exposes both the final results and intermediate results (variables).

    Args:
      code: Python code to run.
      global_vars: An optional dict of
      permission: Permission for the Python code to run.
      returns_stdout: If True, the stdout (a str) will be returned.
      outputs_intermediate: Applicable when returns_stdout is False. If True,
        intermediate output will be outputted as a dict, with the last line's
        value accessible by key '__result__' and the std output accessible by
        key '__stdout__'. Otherwise the value of the last line will be returned.
      sandbox: If True, run code in sandbox; If False, run code in current
        process. If None, run in sandbox first, if the output could not be
        serialized and pass to current process, run the code again in current
        process.
      timeout: Execution timeout in seconds. If None, wait the code the complete.

    Returns:
      The value of the last line of the code block. Or a dict of variable
      names of all locals to their evaluated values as the output of the code to
      run. The value for the last line can be accessed by key '__result__'. Or the
      stdout as a str.

    Raises:
      TimeoutError: If the execution time exceeds the timeout.
      Exception: Exception  that are raised from the code.
    """
    return maybe_sandbox_call(
        evaluate,
        code=code,
        global_vars=global_vars,
        permission=permission,
        returns_stdout=returns_stdout,
        outputs_intermediate=outputs_intermediate,
        sandbox=sandbox,
        timeout=timeout,
    )

sandbox_call

sandbox_call(
    func: Callable[..., Any],
    *args: Any,
    timeout: float | None = None,
    **kwargs: Any
) -> Any

Calls a function with sandboxing.

Parameters:

Name	Type	Description	Default
func	Callable[..., Any]	Function to call.	required
*args	Any	Positional arguments for func	()
timeout	float | None	Execution timeout in seconds. If None, wait func to complete.	None
**kwargs	Any	Keyword arguments for func.	{}

Returns:

Type	Description
Any	Return value from func.

Raises:

Type	Description
TimeoutError	If the execution time exceeds the timeout.
SerializationError	If func (or its arguments) cannot be shipped to the sandbox subprocess.
Exception	Exception raised from func.

Source code in pygx/coding/_execution.py

def sandbox_call(
    func: Callable[..., Any],
    *args: Any,
    timeout: float | None = None,
    **kwargs: Any,
) -> Any:
    """Calls a function with sandboxing.

    Args:
      func: Function to call.
      *args: Positional arguments for `func`
      timeout: Execution timeout in seconds. If None, wait `func` to complete.
      **kwargs: Keyword arguments for `func`.

    Returns:
      Return value from `func`.

    Raises:
      TimeoutError: If the execution time exceeds the timeout.
      SerializationError: If `func` (or its arguments) cannot be shipped to the
        sandbox subprocess.
      Exception: Exception raised from `func`.
    """
    # Serialize the payload up front (in the parent) so a non-picklable target
    # surfaces as a `SerializationError` here rather than as an opaque child
    # crash, and so `maybe_sandbox_call`'s automatic fallback can catch it.
    payload = _dumps_payload((func, args, kwargs, get_context()))

    ctx = multiprocessing.get_context(sandbox_start_method())
    q = ctx.Queue()
    # `get_context(str)` is typed to return the loosely-typed `BaseContext`
    # (only literal method names resolve to a concrete context exposing
    # `Process`); the attribute exists on every concrete context at runtime.
    p = ctx.Process(  # pyright: ignore[reportAttributeAccessIssue]
        target=_sandbox_target, args=(payload, q)
    )
    try:
        p.start()
        x = q.get(timeout=timeout)
    except queue.Empty as e:
        if p.is_alive():
            # We use `kill` instead of `terminate` to release process resources
            # right away.
            p.kill()
        raise TimeoutError(f'Execution time exceed {timeout} seconds.') from e
    finally:
        q.close()

    if isinstance(x, Exception):
        raise x
    try:
        return pickle.loads(x)
    except Exception as e:
        raise errors.SerializationError(
            'Cannot deserialize the output from sandbox.', e
        ) from e

make_function

make_function(
    name: str,
    args: list[str],
    body: list[str],
    *,
    exec_globals: dict[str, Any] | None = None,
    exec_locals: dict[str, Any] | None = None,
    return_type: Any = NO_TYPE_ANNOTATION
)

Creates a function dynamically from source.

Source code in pygx/coding/_function_generation.py

def make_function(
    name: str,
    args: list[str],
    body: list[str],
    *,
    exec_globals: dict[str, Any] | None = None,
    exec_locals: dict[str, Any] | None = None,
    return_type: Any = NO_TYPE_ANNOTATION,
):
    """Creates a function dynamically from source."""
    if exec_locals is None:
        exec_locals = {}
    if return_type != NO_TYPE_ANNOTATION:
        exec_locals['_return_type'] = return_type
        return_annotation = '->_return_type'
    else:
        return_annotation = ''
    args_str = ', '.join(args)
    body_str = '\n'.join(f'  {line}' for line in body)
    fn_def = f' def {name}({args_str}){return_annotation}:\n{body_str}'
    local_vars = ', '.join(exec_locals.keys())
    fn_def = f'def _make_fn({local_vars}):\n{fn_def}\n return {name}'
    ns: dict[str, Any] = {}
    exec(fn_def, exec_globals, ns)  # pylint: disable=exec-used
    return ns['_make_fn'](**exec_locals)

get_permission

get_permission() -> CodePermission | None

Gets the current permission for code execution.

Source code in pygx/coding/_permissions.py

def get_permission() -> CodePermission | None:
    """Gets the current permission for code execution."""
    return concurrent.thread_local_get(_TLS_CODE_RUN_PERMISSION, None)

permission

permission(perm: CodePermission) -> Iterator[CodePermission]

Context manager for controling the permission for code execution.

When the permission context manager is nested, the outtermost permission will be used. This design allows users to control permission at the top level.

Parameters:

Name	Type	Description	Default
perm	CodePermission	Code execution permission.	required

Yields:

Type	Description
CodePermission	Actual permission applied.

Source code in pygx/coding/_permissions.py

@contextlib.contextmanager
def permission(perm: CodePermission) -> Iterator[CodePermission]:
    """Context manager for controling the permission for code execution.

    When the `permission` context manager is nested, the outtermost permission
    will be used. This design allows users to control permission at the top level.

    Args:
      perm: Code execution permission.

    Yields:
      Actual permission applied.
    """

    outter_perm = concurrent.thread_local_get(_TLS_CODE_RUN_PERMISSION, None)

    # Use the top-level permission as the actual permission
    if outter_perm is not None:
        perm = outter_perm

    concurrent.thread_local_set(_TLS_CODE_RUN_PERMISSION, perm)

    try:
        yield perm
    finally:
        if outter_perm is None:
            concurrent.thread_local_del(_TLS_CODE_RUN_PERMISSION)

options: show_root_heading: false show_root_toc_entry: false members_order: source heading_level: 2