import contextlib from abc import ABC, abstractmethod from typing import Any, Dict, List, Literal, Optional, Tuple LogErrorLevel = Literal["info", "warning", "error"] ArtifactType = Literal["chart", "image", "df", "file", "txt", "svg", "html"] class PluginContext(ABC): """ interface for API to interact with execution environment of plugin The runtime will provide an implementation of this interface to the plugin. Plugin could use the API provded withotu need to implement this interface. """ @property @abstractmethod def env_id(self) -> str: """get the environment id of the plugin""" ... @property @abstractmethod def session_id(self) -> str: """get the session id of the plugin""" ... @property @abstractmethod def execution_id(self) -> str: """get the execution id of the plugin""" ... @abstractmethod def add_artifact( self, name: str, file_name: str, type: ArtifactType, val: Any, desc: Optional[str] = None, ) -> str: """ add an artifact to the execution context :param name: the name of the artifact :param file_name: the name of the file :param type: the type of the artifact :param val: the value of the artifact :param desc: the description of the artifact :return: the id of the artifact """ ... @abstractmethod def create_artifact_path( self, name: str, file_name: str, type: ArtifactType, desc: str, ) -> Tuple[str, str]: """ create a path for an artifact and the plugin can use this path to save the artifact. This methods is provided for the plugin to save the artifact by itself rather than saving by the runtime, for general cases when the file content could be passed directly, the plugin should use add_artifact instead. :param name: the name of the artifact :param file_name: the name of the file :param type: the type of the artifact :param desc: the description of the artifact :return: the id and the path of the artifact """ ... @abstractmethod def get_session_var( self, variable_name: str, default: Optional[str], ) -> Optional[str]: """ get a session variable from the context :param variable_name: the name of the variable :param default: the default value of the variable :return: the value of the variable """ ... @abstractmethod def log(self, level: LogErrorLevel, tag: str, message: str) -> None: """log a message from the plugin""" @abstractmethod def get_env(self, plugin_name: str, variable_name: str) -> str: """get an environment variable from the context""" class TestPluginContxt(PluginContext): """ This plugin context is used for testing purpose. """ def __init__(self, temp_dir: str) -> None: self._session_id = "test" self._env_id = "test" self._execution_id = "test" self._logs: List[Tuple[LogErrorLevel, str, str]] = [] self._env: Dict[str, str] = {} self._session_var: Dict[str, str] = {} self._temp_dir = temp_dir self._artifacts: List[Dict[str, str]] = [] @property def env_id(self) -> str: return "test" @property def session_id(self) -> str: return "test" @property def execution_id(self) -> str: return "test" def add_artifact( self, name: str, file_name: str, type: ArtifactType, val: Any, desc: Optional[str] = None, ) -> str: id = f"test_artifact_id_{len(self._artifacts)}" self._artifacts.append( { "id": id, "name": name, "file_name": file_name, "type": type, "desc": desc or "", }, ) return id def create_artifact_path( self, name: str, file_name: str, type: ArtifactType, desc: str, ) -> Tuple[str, str]: id = f"test_artifact_id_{len(self._artifacts)}" self._artifacts.append( { "id": id, "name": name, "file_name": file_name, "type": type, "desc": desc or "", }, ) return id, self._temp_dir + "/" + file_name def log(self, level: LogErrorLevel, tag: str, message: str) -> None: return self._logs.append((level, tag, message)) def get_env(self, plugin_name: str, variable_name: str) -> str: return self._env[plugin_name + "_" + variable_name] def get_session_var( self, variable_name: str, default: Optional[str], ) -> Optional[str]: return self._session_var.get(variable_name, default) @contextlib.contextmanager def temp_context(workspace_dir: Optional[str] = None): import os import shutil import tempfile import uuid if workspace_dir is None: workspace_dir = tempfile.mkdtemp() else: workspace_dir = os.path.join(workspace_dir, str(uuid.uuid4())) os.makedirs(workspace_dir) try: yield TestPluginContxt(workspace_dir) finally: shutil.rmtree(workspace_dir)