API Reference¶

def __enter__(self) -> "ArtifactSandboxSession":
    """Enter the context manager."""
    self._session.__enter__()
    return self

def __exit__(
    self,
    exc_type: type[BaseException] | None,
    exc_val: BaseException | None,
    exc_tb: TracebackType | None,
) -> None:
    """Exit the context manager."""
    return self._session.__exit__(exc_type, exc_val, exc_tb)

def __getattr__(self, name: str) -> Any:
    """Delegate any other attributes/methods to the underlying session."""
    return getattr(self._session, name)

def run(self, code: str, libraries: list | None = None) -> ExecutionResult:
    """Run code in the sandbox session and extract any generated artifacts.

    This method executes the provided code in an isolated environment and captures any
    generated artifacts (e.g., plots, figures). When plotting is enabled, it delegates
    to the language handler's run_with_artifacts method for language-specific artifact
    extraction.

    Args:
        code (str): The code to execute. Can include plotting commands from matplotlib,
                    seaborn, plotly, or other visualization libraries.
        libraries (list | None, optional): Additional libraries to install before running
                                            the code. Defaults to None.

    Returns:
        ExecutionResult: An object containing:
            - exit_code (int): The exit code of the execution
            - stdout (str): Standard output from the code execution
            - stderr (str): Standard error from the code execution
            - plots (list[Plot]): List of captured plots, each containing:
                - content_base64 (str): Base64 encoded plot data
                - format (PlotFormat): Format of the plot (e.g., 'png', 'svg')

    Raises:
        LanguageNotSupportPlotError: If the language does not support plot detection

    Examples:
        Basic plotting example:
        ```python
        with ArtifactSandboxSession(
            lang="python",
            verbose=True,
            image="ghcr.io/vndee/sandbox-python-311-bullseye"
        ) as session:
            code = '''
            import matplotlib.pyplot as plt
            import numpy as np

            x = np.linspace(0, 10, 100)
            y = np.sin(x)
            plt.plot(x, y)
            plt.title('Sine Wave')
            plt.show()
            '''
            result = session.run(code)
            print(f"Generated {len(result.plots)} plots")
        ```

        Multiple plot types and libraries:
        ```python
        code = '''
        import matplotlib.pyplot as plt
        import seaborn as sns
        import plotly.express as px
        import pandas as pd
        import numpy as np

        # Matplotlib plot
        plt.figure(figsize=(10, 6))
        x = np.linspace(0, 10, 100)
        plt.plot(x, np.sin(x))
        plt.title('Matplotlib: Sine Wave')
        plt.show()

        # Seaborn plot
        data = pd.DataFrame({
            'x': np.random.randn(100),
            'y': np.random.randn(100)
        })
        sns.scatterplot(data=data, x='x', y='y')
        plt.title('Seaborn: Scatter Plot')
        plt.show()

        # Plotly plot
        fig = px.line(data, x='x', y='y', title='Plotly: Line Plot')
        fig.show()
        '''

        result = session.run(code, libraries=['plotly'])

        # Save the generated plots
        for i, plot in enumerate(result.plots):
            with open(f'plot_{i}.{plot.format.value}', 'wb') as f:
                f.write(base64.b64decode(plot.content_base64))
        ```

        Installing additional libraries:
        ```python
        code = '''
        import torch
        import torch.nn as nn
        print(f"PyTorch version: {torch.__version__}")
        '''

        result = session.run(code, libraries=['torch'])
        print(result.stdout)
        ```

    """
    # Check if plotting is enabled and language supports it
    if self.enable_plotting and not self._session.language_handler.is_support_plot_detection:
        raise LanguageNotSupportPlotError(self._session.language_handler.name)

    # Delegate to language handler for language-specific artifact extraction
    result, plots = self._session.language_handler.run_with_artifacts(
        container=self._session,  # type: ignore[arg-type]
        code=code,
        libraries=libraries,
        enable_plotting=self.enable_plotting,
        output_dir="/tmp/sandbox_plots",
    )

    return ExecutionResult(
        exit_code=result.exit_code,
        stdout=result.stdout,
        stderr=result.stderr,
        plots=plots,
    )

def success(self) -> bool:
    r"""Check if the execution was successful (exit code is 0).

    Returns:
        bool: True if `exit_code` is 0, False otherwise.

    """
    return not self.exit_code

def text(self) -> str:
    r"""Get the text representation of the console output (stdout).

    .. deprecated:: 0.1.0
        The `text` property is deprecated and will be removed in a future version.
        Use the `stdout` attribute directly instead.

    Returns:
        str: The content of the standard output stream.

    """
    warnings.warn(
        "The 'text' property is deprecated and will be removed in a future version. "
        "Use 'stdout' attribute directly instead.",
        DeprecationWarning,
        stacklevel=2,
    )
    return self.stdout

def add_pattern(self, pattern: SecurityPattern) -> None:
    """Add a security pattern to the policy."""
    if self.patterns is None:
        self.patterns = []
    self.patterns.append(pattern)

def add_restricted_module(self, module: RestrictedModule) -> None:
    """Add a restricted module to the policy."""
    if self.restricted_modules is None:
        self.restricted_modules = []
    self.restricted_modules.append(module)

@field_validator("pattern")
@classmethod
def validate_pattern(cls, v: str) -> str:
    """Validate that the pattern is a valid regex pattern."""
    try:
        re.compile(v)
    except re.error as e:
        raise InvalidRegexPatternError(v) from e
    return v

def extract_plots(self, container: "ContainerProtocol", output_dir: str) -> list[PlotOutput]:  # noqa: ARG002
    """Extract plots from the code.

    Subclasses should override this method to provide custom plot extraction logic.

    Args:
        container: The container protocol instance to run code in
        output_dir: Directory where plots should be saved

    Returns:
        list[PlotOutput]: List of plot outputs

    """
    return []

def filter_comments(self, code: str) -> str:
    """Filter out comments from code in a language-specific way.

    Args:
        code (str): The code to filter comments from.

    Returns:
        str: The code with comments removed.

    """
    # First remove multi-line comments
    code = re.sub(self.get_multiline_comment_patterns(), "", code)

    # Then handle single-line comments
    filtered_lines = []
    for line in code.split("\n"):
        # Remove inline comments
        clean_line = re.sub(self.get_inline_comment_patterns(), "", line)
        # Keep the line if it has non-whitespace content
        if clean_line.strip():
            filtered_lines.append(clean_line)
        else:
            # Preserve empty lines for readability
            filtered_lines.append("")
    return "\n".join(filtered_lines)

def get_execution_commands(self, code_file: str) -> list[str]:
    """Get commands to execute code file."""
    if not self.config.execution_commands:
        raise CommandFailedError(self.config.name, 1, "No execution commands found")
    return [command.format(file=code_file) for command in self.config.execution_commands]

@abstractmethod
def get_import_patterns(self, module: str) -> str:
    """Get the regex patterns for import statements."""

@staticmethod
@abstractmethod
def get_inline_comment_patterns() -> str:
    """Get the regex for inline comment patterns."""

def get_library_installation_command(self, library: str) -> str:
    """Get command to install library."""
    if not self.config.package_manager:
        raise PackageManagerError(self.config.name)
    return f"{self.config.package_manager} {library}"

@staticmethod
@abstractmethod
def get_multiline_comment_patterns() -> str:
    """Get the regex patterns for multiline comment."""

def inject_plot_detection_code(self, code: str) -> str:
    """Inject code to detect and capture plots.

    Subclasses should override this method to provide custom plot detection code.

    Args:
        code: The code to inject plot detection code into.

    Returns:
        The code with plot detection code injected.

    """
    return code

def run_with_artifacts(
    self,
    container: "ContainerProtocol",
    code: str,
    libraries: list | None = None,
    enable_plotting: bool = True,
    output_dir: str = "/tmp/sandbox_plots",
) -> tuple[Any, list[PlotOutput]]:
    """Run code and extract artifacts (plots) in a language-specific manner.

    This method provides a language-specific implementation for running code
    with artifact extraction. Languages that support plot detection can override
    this method to provide custom artifact extraction logic.

    Args:
        container: The container protocol instance to run code in
        code: The code to execute
        libraries: Optional list of libraries to install before running
        enable_plotting: Whether to enable plot detection and extraction
        output_dir: Directory where plots should be saved

    Returns:
        tuple: (execution_result, list_of_plots)

    """
    # Default implementation for languages without plot support
    if enable_plotting and self.is_support_plot_detection:
        # Inject plot detection code
        injected_code = self.inject_plot_detection_code(code)

        # Run the code with plot detection
        result = container.run(injected_code, libraries)

        # Extract plots
        plots = self.extract_plots(container, output_dir)

        return result, plots
    # Run code without plot detection
    result = container.run(code, libraries)
    return result, []