diff --git a/core/analysis_scripts/analysis_scripts/__init__.py b/core/analysis_scripts/analysis_scripts/__init__.py index c4a4f0a..2275c45 100644 --- 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 new file mode 100644 index 0000000..bcccf68 --- /dev/null +++ b/core/analysis_scripts/analysis_scripts/env_tool.py @@ -0,0 +1,132 @@ +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) diff --git a/core/analysis_scripts/analysis_scripts/plugins.py b/core/analysis_scripts/analysis_scripts/plugins.py index 420f6e0..7c99cc4 100644 --- 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 new file mode 100644 index 0000000..85a06bb --- /dev/null +++ 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