allow plugins to be installed and run in virtual environments

menzel-r · menzel-r · commit b72a1eebc7dc · 2024-12-13T16:14:36.000-05:00
diff --git a/core/analysis_scripts/analysis_scripts/__init__.py b/core/analysis_scripts/analysis_scripts/__init__.py
@@ -1,3 +1,4 @@
 from .base_class import AnalysisScript
-from .plugins import available_plugins, find_plugins, plugin_requirements, \
+from .env_tool import VirtualEnvManager
+from .plugins import available_plugins, plugin_requirements, \
                      run_plugin, UnknownPluginError
diff --git a/core/analysis_scripts/analysis_scripts/env_tool.py b/core/analysis_scripts/analysis_scripts/env_tool.py
@@ -0,0 +1,149 @@
+from pathlib import Path
+from subprocess import CalledProcessError, PIPE, run, STDOUT
+from tempfile import TemporaryDirectory
+import venv
+
+
+def _process_output(output):
+    """Converts bytes string to list of String lines.
+
+    Args:
+        output: Bytes string.
+
+    Returns:
+        List of strings.
+    """
+    return [x for x in output.decode("utf-8").split("\n") if x]
+
+
+class VirtualEnvManager(object):
+    """Helper class for creating/running simple command in a virtual environment."""
+    def __init__(self, path):
+        self.path = Path(path)
+        self.activate = f"source {self.path / 'bin' / 'activate'}"
+
+    @staticmethod
+    def _execute(commands):
+        """Runs input commands through bash in a child process.
+
+        Args:
+            commands: List of string commands.
+
+        Returns:
+            List of string output.
+        """
+        with TemporaryDirectory() as tmp:
+            script_path = Path(tmp) / "script"
+            with open(script_path, "w") as script:
+                script.write("\n".join(commands))
+            try:
+                process = run(["bash", str(script_path)], stdout=PIPE, stderr=STDOUT,
+                              check=True)
+            except CalledProcessError as err:
+                for line in _process_output(err.output):
+                    print(line)
+                raise
+            return _process_output(process.stdout)
+
+    def _execute_python_script(self, commands):
+        """Runs input python code in bash in a child process.
+
+        Args:
+            commands: List of string python code lines.
+
+        Returns:
+            List of string output.
+        """
+        with TemporaryDirectory() as tmp:
+            script_path = Path(tmp) / "python_script"
+            with open(script_path, "w") as script:
+                script.write("\n".join(commands))
+            commands = [self.activate, f"python3 {str(script_path)}"]
+            return self._execute(commands)
+
+    def create_env(self):
+        """Creates the virtual environment."""
+        venv.create(self.path, with_pip=True)
+
+    def destroy_env(self):
+        """Destroys the virtual environment."""
+        raise NotImplementedError("this feature is not implemented yet.")
+
+    def install_package(self, name):
+        """Installs a package in the virtual environment.
+
+        Args:
+            name: String name of the package.
+
+        Returns:
+            List of string output.
+        """
+        commands = [self.activate, "python3 -m pip --upgrade pip",
+                    f"python3 -m pip install {name}"]
+        return self._execute(commands)
+
+    def list_plugins(self):
+        """Returns a list of plugins that are available in the virtual environment.
+
+        Returns:
+            List of plugins.
+        """
+        python_script = [
+            "from analysis_scripts import available_plugins",
+            "for plugin in available_plugins():",
+            "    print(plugin)"
+        ]
+        return self._execute_python_script(python_script)
+
+    def run_analysis_plugin(self, name, catalog, output_directory, config=None):
+        """Returns a list of paths to figures created by the plugin from the virtual
+           environment.
+
+        Args:
+             name: String name of the analysis package.
+             catalog: Path to the data catalog.
+             output_directory: Path to the output directory.
+
+        Returns:
+            List of figure paths.
+        """
+        if config:
+            python_script = [f"config = {str(config)}",]
+        else:
+            python_script = ["config = None",]
+        python_script += [
+            "from analysis_scripts import run_plugin",
+            f"paths = run_plugin('{name}', '{catalog}', '{output_directory}', config=config)",
+            "for path in paths:",
+            "    print(path)"
+        ]
+        return self._execute_python_script(python_script)
+
+    def uninstall_package(self, name):
+        """Uninstalls a package from the virtual environment.
+
+        Args:
+            name: String name of the package.
+
+        Returns:
+            List of string output.
+        """
+        commands = [self.activate, f"pip uninstall {name}"]
+        return self._execute(commands)
+
+
+def analysis_script_test():
+    env = VirtualEnvManager("foo")
+    name = "freanalysis_clouds"
+    url = "https://github.com/noaa-gfdl/analysis-scripts.git"
+    with TemporaryDirectory() as tmp:
+        tmp_path = Path(tmp)
+        run(["git", "clone", url, str(tmp_path / "scripts")])
+        output = env.install_package(str(tmp_path / "scripts" / "core" / "figure_tools"))
+        output = env.install_package(str(tmp_path / "scripts" / "core" / "analysis_scripts"))
+        output = env.install_package(str(tmp_path / "scripts" / "user-analysis-scripts" / name))
+    env.run_analysis_plugin(name, "fake-catalog", ".", config={"a": "b"})
+
+
+if __name__ == "__main__":
+    analysis_script_test()
diff --git a/core/analysis_scripts/analysis_scripts/plugins.py b/core/analysis_scripts/analysis_scripts/plugins.py
@@ -1,57 +1,15 @@
 import importlib
 import inspect
 import pkgutil
-import sys
 
 from .base_class import AnalysisScript
 
 
-class _PathAdjuster(object):
-    """Helper class to adjust where python tries to import modules from."""
-    def __init__(self, path):
-        """Initialize the object.
-
-        Args:
-            path: Path to look in for python modules and packages.
-        """
-        self.path = path
-        self.old_sys_path = sys.path
-
-    def __enter__(self):
-        """Adjusts the sys path so the modules and packages can be imported."""
-        if self.path not in sys.path:
-            sys.path.insert(0, self.path)
-        return self
-
-    def __exit__(self, exception_type, exception_value, traceback):
-        """Undoes the sys path adjustment."""
-        if sys.path != self.old_sys_path:
-            sys.path = self.old_sys_path
-
-
 # Dictionary of found plugins.
 discovered_plugins = {}
-
-
-def find_plugins(path=None):
-    """Find all installed python modules with names that start with 'freanalysis_'.
-
-    Args:
-        path: Custom directory where modules and packages are installed.
-    """
-    if path:
-        path = [path,]
-    for finder, name, ispkg in pkgutil.iter_modules(path):
-        if name.startswith("freanalysis_"):
-            if path:
-                with _PathAdjuster(path[0]):
-                    discovered_plugins[name] = importlib.import_module(name)
-            else:
-                discovered_plugins[name] = importlib.import_module(name)
-
-
-# Update plugin dictionary.
-find_plugins()
+for finder, name, ispkg in pkgutil.iter_modules():
+    if name.startswith("freanalysis_"):
+        discovered_plugins[name] = importlib.import_module(name)
 
 
 class UnknownPluginError(BaseException):
diff --git a/core/analysis_scripts/tests/test_env_tool.py b/core/analysis_scripts/tests/test_env_tool.py
@@ -0,0 +1,58 @@
+from pathlib import Path
+from platform import python_version_tuple
+from subprocess import CalledProcessError, run
+from tempfile import TemporaryDirectory
+
+from analysis_scripts import VirtualEnvManager
+import pytest
+
+
+def install_helper(tmp):
+    tmp_path = Path(tmp)
+    env_path = tmp_path / "env"
+    env = VirtualEnvManager(env_path)
+    env.create_env()
+    name = "freanalysis_clouds"
+    url = "https://github.com/noaa-gfdl/analysis-scripts.git"
+    run(["git", "clone", url, str(tmp_path / "scripts")])
+    output = env.install_package(str(tmp_path / "scripts" / "core" / "analysis_scripts"))
+    output = env.install_package(str(tmp_path / "scripts" / "core" / "figure_tools"))
+    output = env.install_package(str(tmp_path / "scripts" / "user-analysis-scripts" / name))
+    return tmp_path, env_path, env, name
+
+
+def test_create_env():
+    with TemporaryDirectory() as tmp:
+        env_path = Path(tmp) / "env"
+        env = VirtualEnvManager(env_path)
+        env.create_env()
+        assert env_path.is_dir()
+        test_string = "hello, world"
+        assert env._execute([f'echo "{test_string}"',])[0] == test_string
+
+
+def test_install_plugin():
+    with TemporaryDirectory() as tmp:
+        tmp_path, env_path, env, name = install_helper(tmp)
+        version = ".".join(python_version_tuple()[:2])
+        plugin_path = env_path / "lib" / f"python{version}" / "site-packages" / name
+        assert plugin_path.is_dir()
+
+
+def test_list_plugins():
+    with TemporaryDirectory() as tmp:
+        tmp_path, env_path, env, name = install_helper(tmp)
+        plugins = env.list_plugins()
+        assert plugins[0] == name
+
+
+def test_run_plugin():
+    with TemporaryDirectory() as tmp:
+        tmp_path, env_path, env, name = install_helper(tmp)
+        catalog = tmp_path / "fake-catalog"
+        with pytest.raises(CalledProcessError) as err:
+            env.run_analysis_plugin(name, str(catalog), ".", config={"a": "b"})
+        for line in err._excinfo[1].output.decode("utf-8").split("\n"):
+            if f"No such file or directory: '{str(catalog)}'" in line:
+                return
+        assert False
