poethepoet.context

  1import re
  2from pathlib import Path
  3from typing import TYPE_CHECKING, Any, Dict, Mapping, Optional, Tuple, Union
  4
  5if TYPE_CHECKING:
  6    from .config import PoeConfig
  7    from .env.manager import EnvVarsManager
  8    from .executor import PoeExecutor
  9    from .ui import PoeUi
 10
 11
 12class RunContext:
 13    config: "PoeConfig"
 14    ui: "PoeUi"
 15    env: "EnvVarsManager"
 16    dry: bool
 17    poe_active: Optional[str]
 18    project_dir: Path
 19    multistage: bool = False
 20    exec_cache: Dict[str, Any]
 21    captured_stdout: Dict[Tuple[str, ...], str]
 22
 23    def __init__(
 24        self,
 25        config: "PoeConfig",
 26        ui: "PoeUi",
 27        env: Mapping[str, str],
 28        dry: bool,
 29        poe_active: Optional[str],
 30        multistage: bool = False,
 31        cwd: Optional[Union[Path, str]] = None,
 32    ):
 33        from .env.manager import EnvVarsManager
 34
 35        self.config = config
 36        self.ui = ui
 37        self.project_dir = Path(config.project_dir)
 38        self.dry = dry
 39        self.poe_active = poe_active
 40        self.multistage = multistage
 41        self.exec_cache = {}
 42        self.captured_stdout = {}
 43
 44        # Init root EnvVarsManager
 45        self.env = EnvVarsManager(self.config, self.ui, base_env=env, cwd=cwd)
 46        for config_part in self.config.partitions():
 47            self.env.apply_env_config(
 48                envfile=config_part.get("envfile", None),
 49                config_env=config_part.get("env", None),
 50                config_dir=config_part.config_dir,
 51                config_working_dir=config_part.cwd,
 52            )
 53
 54    def _get_dep_values(
 55        self, used_task_invocations: Mapping[str, Tuple[str, ...]]
 56    ) -> Dict[str, str]:
 57        """
 58        Get env vars from upstream tasks declared via the uses option.
 59        """
 60        return {
 61            var_name: self.get_task_output(invocation)
 62            for var_name, invocation in used_task_invocations.items()
 63        }
 64
 65    def save_task_output(self, invocation: Tuple[str, ...], captured_stdout: bytes):
 66        """
 67        Store the stdout data from a task so that it can be reused by other tasks
 68        """
 69        try:
 70            self.captured_stdout[invocation] = captured_stdout.decode()
 71        except UnicodeDecodeError:
 72            # Attempt to recover in case a specific encoding is configured
 73            io_encoding = self.env.get("PYTHONIOENCODING")
 74            if io_encoding:
 75                self.captured_stdout[invocation] = captured_stdout.decode(io_encoding)
 76            else:
 77                raise
 78
 79    def get_task_output(self, invocation: Tuple[str, ...]):
 80        """
 81        Get the stored stdout data from a task so that it can be reused by other tasks
 82
 83        New lines are replaced with whitespace similar to how unquoted command
 84        interpolation works in bash.
 85        """
 86        return re.sub(r"\s+", " ", self.captured_stdout[invocation].strip("\r\n"))
 87
 88    def get_executor(
 89        self,
 90        invocation: Tuple[str, ...],
 91        env: "EnvVarsManager",
 92        working_dir: Path,
 93        executor_config: Optional[Mapping[str, str]] = None,
 94        capture_stdout: Union[str, bool] = False,
 95    ) -> "PoeExecutor":
 96        from .executor import PoeExecutor
 97
 98        if not executor_config:
 99            if self.ui["executor"]:
100                executor_config = {"type": self.ui["executor"]}
101            else:
102                executor_config = self.config.executor
103
104        return PoeExecutor.get(
105            invocation=invocation,
106            context=self,
107            executor_config=executor_config,
108            env=env,
109            working_dir=working_dir,
110            capture_stdout=capture_stdout,
111            dry=self.dry,
112        )