| """Build step class for setting up a Python environment.""" |
| from pathlib import Path |
| import os |
| from typing import Iterable, Optional |
| |
| from helpers import py_env_utils |
| from slave import base_step |
| |
| |
| class PyEnvSetupStep(base_step.BaseStep): |
| """Step for setting up python envs.""" |
| |
| def __init__(self, |
| additional_import_paths: Optional[Iterable[Path]]=None, |
| local_package_project: str=None, |
| local_packages_dir: str=None, |
| py3: bool=False, |
| requirements_project: str=None, |
| requirements_path: str=None, |
| **kwargs): |
| """Creates a PyEnvSetupStep instance. |
| |
| Args: |
| additional_import_paths: Paths to add to sys.path for the environment |
| local_package_project: manifest project to look a python packages on disk |
| local_packages_dir: Relative path to python packages under ~/gcs-mirror |
| py3: if True use python3 |
| requirements_project: manifest project to look for requirements_path |
| requirements_path: path to requirements file under requirements_project |
| gcs_mirror_path: The directory that mirrors the gcs bucket |
| **kwargs: Any additional args to pass to BaseStep. |
| """ |
| super().__init__(name='py_env_setup', **kwargs) |
| self._gcs_dir = self._properties.get( |
| 'gcs_mirror_path', '/workspace/gcs-mirror') |
| self._additional_import_paths = additional_import_paths |
| self._local_package_project = local_package_project |
| self._local_packages_dir = local_packages_dir |
| self._py3 = py3 |
| self._requirements_path = requirements_path |
| self._requirements_project = requirements_project |
| |
| def run(self) -> bool: |
| """Sets up a Python environment. |
| |
| Adds 'virtual_env_root' to the step data. |
| |
| Returns: |
| True iff the environment was set up correctly. |
| """ |
| py_requirements_dir = self.get_project_path(self._requirements_project) \ |
| if self._requirements_project else None |
| local_package_path = self.get_project_path(self._local_package_project) \ |
| if self._local_package_project else None |
| venv_root = py_env_utils.create_virtual_env( |
| self, |
| self.get_tmp_dir(), |
| gcs_dir=self._gcs_dir, |
| local_packages=self._local_packages_dir, |
| py3=self._py3) |
| py_env_utils.add_import_paths(venv_root, [Path.home() / 'depot_tools']) |
| self.add_step_data('virtual_env_root', venv_root) |
| if py_requirements_dir: |
| requirements = os.path.join(py_requirements_dir, |
| self._requirements_path) |
| py_env_utils.install_py_deps( |
| self, venv_root, requirements, self._gcs_dir, self._local_packages_dir |
| ) |
| if local_package_path: |
| py_env_utils.install_package( |
| self, |
| venv_root, |
| local_package_path, |
| self._gcs_dir, |
| local_packages=self._local_packages_dir |
| ) |
| |
| if self._additional_import_paths: |
| py_env_utils.add_import_paths(venv_root, self._additional_import_paths) |
| return True |
