zrb 0.0.117__py3-none-any.whl → 0.0.119__py3-none-any.whl

zrb/task/base_task/component/base_task_model.py

@@ -0,0 +1,258 @@
+from zrb.helper.typing import (
+    Any, Callable, Iterable, List, Mapping, Optional, Union
+)
+from zrb.helper.typecheck import typechecked
+from zrb.config.config import show_time
+from zrb.task.any_task import AnyTask
+from zrb.task.any_task_event_handler import (
+    OnTriggered, OnWaiting, OnSkipped, OnStarted, OnReady, OnRetry, OnFailed
+)
+from zrb.helper.log import logger
+from zrb.helper.accessories.color import colored
+from zrb.task_input.any_input import AnyInput
+from zrb.task_group.group import Group
+from zrb.task_env.env import Env
+from zrb.task_env.env_file import EnvFile
+from zrb.task.base_task.component.common_task_model import CommonTaskModel
+from zrb.task.base_task.component.pid_model import PidModel
+from zrb.task.base_task.component.trackers import TimeTracker
+from zrb.config.config import env_prefix
+from zrb.helper.string.modification import double_quote
+from zrb.helper.string.conversion import to_variable_name
+
+import datetime
+import os
+import sys
+
+LOG_NAME_LENGTH = 20
+
+
+@typechecked
+class BaseTaskModel(CommonTaskModel, PidModel, TimeTracker):
+    def __init__(
+        self,
+        name: str,
+        group: Optional[Group] = None,
+        description: str = '',
+        inputs: List[AnyInput] = [],
+        envs: Iterable[Env] = [],
+        env_files: Iterable[EnvFile] = [],
+        icon: Optional[str] = None,
+        color: Optional[str] = None,
+        retry: int = 2,
+        retry_interval: Union[int, float] = 1,
+        upstreams: Iterable[AnyTask] = [],
+        checkers: Iterable[AnyTask] = [],
+        checking_interval: Union[int, float] = 0,
+        run: Optional[Callable[..., Any]] = None,
+        on_triggered: Optional[OnTriggered] = None,
+        on_waiting: Optional[OnWaiting] = None,
+        on_skipped: Optional[OnSkipped] = None,
+        on_started: Optional[OnStarted] = None,
+        on_ready: Optional[OnReady] = None,
+        on_retry: Optional[OnRetry] = None,
+        on_failed: Optional[OnFailed] = None,
+        should_execute: Union[bool, str, Callable[..., bool]] = True,
+        return_upstream_result: bool = False
+    ):
+        self.__rjust_full_cmd_name: Optional[str] = None
+        self.__has_cli_interface = False
+        self.__complete_name: Optional[str] = None
+        CommonTaskModel.__init__(
+            self,
+            name=name,
+            group=group,
+            description=description,
+            inputs=inputs,
+            envs=envs,
+            env_files=env_files,
+            icon=icon,
+            color=color,
+            retry=retry,
+            retry_interval=retry_interval,
+            upstreams=upstreams,
+            checkers=checkers,
+            checking_interval=checking_interval,
+            run=run,
+            on_triggered=on_triggered,
+            on_waiting=on_waiting,
+            on_skipped=on_skipped,
+            on_started=on_started,
+            on_ready=on_ready,
+            on_retry=on_retry,
+            on_failed=on_failed,
+            should_execute=should_execute,
+            return_upstream_result=return_upstream_result
+        )
+        PidModel.__init__(self)
+        TimeTracker.__init__(self)
+        self.__args: List[Any] = []
+        self.__kwargs: Mapping[str, Any] = {}
+
+    def _set_args(self, args: Iterable[Any]):
+        self.__args = list(args)
+
+    def _set_kwargs(self, kwargs: Mapping[str, Any]):
+        self.__kwargs = kwargs
+
+    def log_debug(self, message: Any):
+        prefix = self.__get_log_prefix()
+        colored_message = colored(
+            f'{prefix} • {message}', attrs=['dark']
+        )
+        logger.debug(colored_message)
+
+    def log_warn(self, message: Any):
+        prefix = self.__get_log_prefix()
+        colored_message = colored(
+            f'{prefix} • {message}', attrs=['dark']
+        )
+        logger.warning(colored_message)
+
+    def log_info(self, message: Any):
+        prefix = self.__get_log_prefix()
+        colored_message = colored(
+            f'{prefix} • {message}', attrs=['dark']
+        )
+        logger.info(colored_message)
+
+    def log_error(self, message: Any):
+        prefix = self.__get_log_prefix()
+        colored_message = colored(
+            f'{prefix} • {message}', color='red', attrs=['bold']
+        )
+        logger.error(colored_message, exc_info=True)
+
+    def log_critical(self, message: Any):
+        prefix = self.__get_log_prefix()
+        colored_message = colored(
+            f'{prefix} • {message}', color='red', attrs=['bold']
+        )
+        logger.critical(colored_message, exc_info=True)
+
+    def print_out(self, message: Any, trim_message: bool = True):
+        prefix = self.__get_colored_print_prefix()
+        message_str = f'{message}'.rstrip() if trim_message else f'{message}'
+        print(f'🤖 ○ {prefix} • {message_str}', file=sys.stderr)
+        sys.stderr.flush()
+
+    def print_err(self, message: Any, trim_message: bool = True):
+        prefix = self.__get_colored_print_prefix()
+        message_str = f'{message}'.rstrip() if trim_message else f'{message}'
+        print(f'🤖 △ {prefix} • {message_str}', file=sys.stderr)
+        sys.stderr.flush()
+
+    def print_out_dark(self, message: Any, trim_message: bool = True):
+        message_str = f'{message}'
+        self.print_out(colored(message_str, attrs=['dark']), trim_message)
+
+    def _print_result(self, result: Any):
+        if result is None:
+            return
+        if self._return_upstream_result:
+            # if _return_upstream_result, result is list (see: self._run_all)
+            upstreams = self._get_upstreams()
+            upstream_results = list(result)
+            for upstream_index, upstream_result in enumerate(upstream_results):
+                upstreams[upstream_index]._print_result(upstream_result)
+            return
+        self.print_result(result)
+
+    def print_result(self, result: Any):
+        '''
+        Print result to stdout so that it can be processed further.
+        e.g.: echo $(zrb explain solid) > solid-principle.txt
+
+        You need to override this method
+        if you want to show the result differently.
+        '''
+        print(result)
+
+    def _play_bell(self):
+        print('\a', end='', file=sys.stderr, flush=True)
+
+    def _show_done_info(self):
+        elapsed_time = self._get_elapsed_time()
+        self.print_out_dark(f'Completed in {elapsed_time} seconds')
+        self._play_bell()
+
+    def _show_env_prefix(self):
+        if env_prefix == '':
+            return
+        colored_env_prefix = colored(env_prefix, color='yellow')
+        colored_label = colored('Your current environment: ', attrs=['dark'])
+        print(colored(f'{colored_label}{colored_env_prefix}'), file=sys.stderr)
+
+    def _show_run_command(self):
+        params: List[str] = [double_quote(arg) for arg in self.__args]
+        for task_input in self._get_combined_inputs():
+            if task_input.is_hidden():
+                continue
+            key = task_input.get_name()
+            kwarg_key = to_variable_name(key)
+            quoted_value = double_quote(str(self.__kwargs[kwarg_key]))
+            params.append(f'--{key} {quoted_value}')
+        run_cmd = self._get_full_cmd_name()
+        run_cmd_with_param = run_cmd
+        if len(params) > 0:
+            param_str = ' '.join(params)
+            run_cmd_with_param += ' ' + param_str
+        colored_command = colored(run_cmd_with_param, color='yellow')
+        colored_label = colored('To run again: ', attrs=['dark'])
+        print(colored(f'{colored_label}{colored_command}'), file=sys.stderr)
+
+    def __get_colored_print_prefix(self) -> str:
+        return self.__get_colored(self.__get_print_prefix())
+
+    def __get_colored(self, text: str) -> str:
+        return colored(text, color=self.get_color())
+
+    def __get_print_prefix(self) -> str:
+        common_prefix = self.__get_common_prefix(show_time=show_time)
+        icon = self.get_icon()
+        rjust_cmd_name = self.__get_rjust_full_cmd_name()
+        return f'{common_prefix} {icon} {rjust_cmd_name}'
+
+    def __get_log_prefix(self) -> str:
+        common_prefix = self.__get_common_prefix(show_time=False)
+        icon = self.get_icon()
+        filled_name = self.__get_rjust_full_cmd_name()
+        return f'{common_prefix} {icon} {filled_name}'
+
+    def __get_common_prefix(self, show_time: bool) -> str:
+        attempt = self._get_attempt()
+        max_attempt = self._get_max_attempt()
+        pid = str(self._get_task_pid()).rjust(6)
+        if show_time:
+            now = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S.%f')[:-3]
+            return f'◷ {now} ❁ {pid} → {attempt}/{max_attempt}'
+        return f'❁ {pid} → {attempt}/{max_attempt}'
+
+    def __get_rjust_full_cmd_name(self) -> str:
+        if self.__rjust_full_cmd_name is not None:
+            return self.__rjust_full_cmd_name
+        complete_name = self._get_full_cmd_name()
+        self.__rjust_full_cmd_name = complete_name.rjust(LOG_NAME_LENGTH, ' ')
+        return self.__rjust_full_cmd_name
+
+    def __get_executable_name(self) -> str:
+        if len(sys.argv) > 0 and sys.argv[0] != '':
+            return os.path.basename(sys.argv[0])
+        return 'zrb'
+
+    def _get_full_cmd_name(self) -> str:
+        if self.__complete_name is not None:
+            return self.__complete_name
+        executable_prefix = ''
+        if self.__has_cli_interface:
+            executable_prefix += self.__get_executable_name() + ' '
+        cmd_name = self.get_cmd_name()
+        if self._group is None:
+            self.__complete_name = f'{executable_prefix}{cmd_name}'
+            return self.__complete_name
+        group_cmd_name = self._group.get_complete_name()
+        self.__complete_name = f'{executable_prefix}{group_cmd_name} {cmd_name}'  # noqa
+        return self.__complete_name
+
+    def _set_has_cli_interface(self):
+        self.__has_cli_interface = True

zrb/task/base_task/component/common_task_model.py

@@ -0,0 +1,282 @@
+from zrb.helper.typing import (
+    Any, Callable, Iterable, List, Mapping, Optional, Union
+)
+from zrb.helper.typecheck import typechecked
+from zrb.task.any_task_event_handler import (
+    OnTriggered, OnWaiting, OnSkipped, OnStarted, OnReady, OnRetry, OnFailed
+)
+from zrb.task.any_task import AnyTask
+from zrb.helper.string.conversion import to_cmd_name
+from zrb.helper.accessories.color import get_random_color
+from zrb.helper.accessories.icon import get_random_icon
+from zrb.helper.util import coalesce_str
+from zrb.task_input.any_input import AnyInput
+from zrb.task_group.group import Group
+from zrb.task_env.constant import RESERVED_ENV_NAMES
+from zrb.task_env.env import Env
+from zrb.task_env.env_file import EnvFile
+
+import os
+
+
+@typechecked
+class CommonTaskModel():
+    def __init__(
+        self,
+        name: str,
+        group: Optional[Group] = None,
+        description: str = '',
+        inputs: List[AnyInput] = [],
+        envs: Iterable[Env] = [],
+        env_files: Iterable[EnvFile] = [],
+        icon: Optional[str] = None,
+        color: Optional[str] = None,
+        retry: int = 2,
+        retry_interval: Union[float, int] = 1,
+        upstreams: Iterable[AnyTask] = [],
+        checkers: Iterable[AnyTask] = [],
+        checking_interval: Union[float, int] = 0,
+        run: Optional[Callable[..., Any]] = None,
+        on_triggered: Optional[OnTriggered] = None,
+        on_waiting: Optional[OnWaiting] = None,
+        on_skipped: Optional[OnSkipped] = None,
+        on_started: Optional[OnStarted] = None,
+        on_ready: Optional[OnReady] = None,
+        on_retry: Optional[OnRetry] = None,
+        on_failed: Optional[OnFailed] = None,
+        should_execute: Union[bool, str, Callable[..., bool]] = True,
+        return_upstream_result: bool = False
+    ):
+        self._name = name
+        self._group = group
+        if group is not None:
+            group.add_task(self)
+        self._description = coalesce_str(description, name)
+        self._inputs = inputs
+        self._envs = envs
+        self._env_files = env_files
+        self._icon = coalesce_str(icon, get_random_icon())
+        self._color = coalesce_str(color, get_random_color())
+        self._retry = retry
+        self._retry_interval = retry_interval
+        self._upstreams = upstreams
+        self._checkers = checkers
+        self._checking_interval = checking_interval
+        self._run_function: Optional[Callable[..., Any]] = run
+        self._on_triggered = on_triggered
+        self._on_waiting = on_waiting
+        self._on_skipped = on_skipped
+        self._on_started = on_started
+        self._on_ready = on_ready
+        self._on_retry = on_retry
+        self._on_failed = on_failed
+        self._should_execute = should_execute
+        self._return_upstream_result = return_upstream_result
+        self.__execution_id = ''
+        self.__allow_add_envs = True
+        self.__allow_add_env_files = True
+        self.__allow_add_inputs = True
+        self.__allow_add_upstreams: bool = True
+        self.__has_already_inject_env_files: bool = False
+        self.__has_already_inject_envs: bool = False
+        self.__has_already_inject_inputs: bool = False
+        self.__has_already_inject_checkers: bool = False
+        self.__has_already_inject_upstreams: bool = False
+        self.__all_inputs: Optional[List[AnyInput]] = None
+
+    def _lock_upstreams(self):
+        self.__allow_add_upstreams = False
+
+    def _set_execution_id(self, execution_id: str):
+        if self.__execution_id == '':
+            self.__execution_id = execution_id
+
+    def _propagate_execution_id(self):
+        execution_id = self.get_execution_id()
+        for upstream_task in self._get_upstreams():
+            upstream_task._set_execution_id(execution_id)
+            upstream_task._propagate_execution_id()
+        for checker_task in self._get_checkers():
+            checker_task._set_execution_id(execution_id)
+            checker_task._propagate_execution_id()
+
+    def get_execution_id(self) -> str:
+        return self.__execution_id
+
+    def set_name(self, new_name: str):
+        if self._description == self._name:
+            self._description = new_name
+        self._name = new_name
+
+    def get_cmd_name(self) -> str:
+        return to_cmd_name(self._name)
+
+    def set_description(self, new_description: str):
+        self._description = new_description
+
+    def get_description(self) -> str:
+        return self._description
+
+    def set_icon(self, new_icon: str):
+        self._icon = new_icon
+
+    def set_color(self, new_color: str):
+        self._color = new_color
+
+    def set_retry(self, new_retry: int):
+        self._retry = new_retry
+
+    def set_should_execute(
+        self, should_execute: Union[bool, str, Callable[..., bool]]
+    ):
+        self._should_execute = should_execute
+
+    def set_retry_interval(self, new_retry_interval: Union[float, int]):
+        self._retry_interval = new_retry_interval
+
+    def set_checking_interval(self, new_checking_interval: Union[float, int]):
+        self._checking_interval = new_checking_interval
+
+    def insert_input(self, *inputs: AnyInput):
+        if not self.__allow_add_inputs:
+            raise Exception(f'Cannot insert inputs for `{self._name}`')
+        self._inputs = list(inputs) + list(self._inputs)
+
+    def add_input(self, *inputs: AnyInput):
+        if not self.__allow_add_inputs:
+            raise Exception(f'Cannot add inputs for `{self._name}`')
+        self._inputs = list(self._inputs) + list(inputs)
+
+    def inject_inputs(self):
+        pass
+
+    def _get_inputs(self) -> List[AnyInput]:
+        if not self.__has_already_inject_inputs:
+            self.inject_inputs()
+            self.__has_already_inject_inputs = True
+        return list(self._inputs)
+
+    def _get_combined_inputs(self) -> Iterable[AnyInput]:
+        ''''
+        Getting all inputs of this task and all its upstream, non-duplicated.
+        '''
+        if self.__all_inputs is not None:
+            return self.__all_inputs
+        self.__all_inputs: List[AnyInput] = []
+        existing_input_names: Mapping[str, bool] = {}
+        # Add task inputs
+        inputs = self._get_inputs()
+        for input_index, first_occurence_task_input in enumerate(inputs):
+            input_name = first_occurence_task_input.get_name()
+            if input_name in existing_input_names:
+                continue
+            # Look for all input with the same name in the current task
+            task_inputs = [
+                candidate
+                for candidate in inputs[input_index:]
+                if candidate.get_name() == input_name
+            ]
+            # Get the last input, and add it to _all_inputs
+            task_input = task_inputs[-1]
+            self.__all_inputs.append(task_input)
+            existing_input_names[input_name] = True
+        # Add upstream inputs
+        for upstream in self._get_upstreams():
+            upstream_inputs = upstream._get_combined_inputs()
+            for upstream_input in upstream_inputs:
+                if upstream_input.get_name() in existing_input_names:
+                    continue
+                self.__all_inputs.append(upstream_input)
+                existing_input_names[upstream_input.get_name()] = True
+        self._lock_upstreams()
+        self.__allow_add_inputs = False
+        return self.__all_inputs
+
+    def insert_env(self, *envs: Env):
+        if not self.__allow_add_envs:
+            raise Exception(f'Cannot insert envs to `{self._name}`')
+        self._envs = list(envs) + list(self._envs)
+
+    def add_env(self, *envs: Env):
+        if not self.__allow_add_envs:
+            raise Exception(f'Cannot add envs to `{self._name}`')
+        self._envs = list(self._envs) + list(envs)
+
+    def inject_envs(self):
+        pass
+
+    def _get_envs(self) -> List[Env]:
+        if not self.__has_already_inject_envs:
+            self.inject_envs()
+            self.__has_already_inject_envs = True
+        return list(self._envs)
+
+    def _get_combined_env(self) -> Mapping[str, Env]:
+        all_envs: Mapping[str, Env] = {}
+        for env_name in os.environ:
+            if env_name in RESERVED_ENV_NAMES:
+                continue
+            all_envs[env_name] = Env(
+                name=env_name, os_name=env_name, should_render=False
+            )
+        for env_file in self._get_env_files():
+            for env in env_file.get_envs():
+                all_envs[env.get_name()] = env
+        for env in self._get_envs():
+            all_envs[env.get_name()] = env
+        self.__allow_add_envs = False
+        self.__allow_add_env_files = False
+        return all_envs
+
+    def insert_env_file(self, *env_files: EnvFile):
+        if not self.__allow_add_env_files:
+            raise Exception(f'Cannot insert env_files to `{self._name}`')
+        self._env_files = list(env_files) + list(self._env_files)
+
+    def add_env_file(self, *env_files: EnvFile):
+        if not self.__allow_add_env_files:
+            raise Exception(f'Cannot add env_files to `{self._name}`')
+        self._env_files = list(self._env_files) + list(env_files)
+
+    def inject_env_files(self):
+        pass
+
+    def insert_upstream(self, *upstreams: AnyTask):
+        if not self.__allow_add_upstreams:
+            raise Exception(f'Cannot insert upstreams to `{self._name}`')
+        self._upstreams = list(upstreams) + list(self._upstreams)
+
+    def add_upstream(self, *upstreams: AnyTask):
+        if not self.__allow_add_upstreams:
+            raise Exception(f'Cannot add upstreams to `{self._name}`')
+        self._upstreams = list(self._upstreams) + list(upstreams)
+
+    def inject_upstreams(self):
+        pass
+
+    def _get_upstreams(self) -> List[AnyTask]:
+        if not self.__has_already_inject_upstreams:
+            self.inject_upstreams()
+            self.__has_already_inject_upstreams = True
+        return list(self._upstreams)
+
+    def get_icon(self) -> str:
+        return self._icon
+
+    def get_color(self) -> str:
+        return self._color
+
+    def _get_env_files(self) -> List[EnvFile]:
+        if not self.__has_already_inject_env_files:
+            self.inject_env_files()
+            self.__has_already_inject_env_files = True
+        return self._env_files
+
+    def inject_checkers(self):
+        pass
+
+    def _get_checkers(self) -> List[AnyTask]:
+        if not self.__has_already_inject_checkers:
+            self.inject_checkers()
+            self.__has_already_inject_checkers = True
+        return list(self._checkers)

zrb/task/base_task/component/pid_model.py

@@ -0,0 +1,17 @@
+from zrb.helper.typecheck import typechecked
+
+import os
+
+
+@typechecked
+class PidModel():
+
+    def __init__(self):
+        self.__task_pid: int = os.getpid()
+
+    def _set_task_pid(self, pid: int):
+        self.__task_pid = pid
+
+    def _get_task_pid(self) -> int:
+        return self.__task_pid
+

zrb/task/base_task/component/renderer.py

@@ -0,0 +1,119 @@
+from zrb.helper.typing import Any, Mapping, Optional, Union
+from zrb.helper.typecheck import typechecked
+from zrb.helper.string.conversion import to_boolean
+from zrb.helper.string.jinja import is_probably_jinja
+from zrb.helper.render_data import DEFAULT_RENDER_DATA
+
+import os
+import jinja2
+
+
+class AnyExtensionFileSystemLoader(jinja2.FileSystemLoader):
+    def get_source(self, environment, template):
+        for search_dir in self.searchpath:
+            file_path = os.path.join(search_dir, template)
+            if os.path.exists(file_path):
+                with open(file_path, 'r') as file:
+                    contents = file.read()
+                return contents, file_path, lambda: False
+        raise jinja2.TemplateNotFound(template)
+
+
+@typechecked
+class Renderer():
+
+    def __init__(self):
+        self.__input_map: Mapping[str, Any] = {}
+        self.__env_map: Mapping[str, str] = {}
+        self.__render_data: Optional[Mapping[str, Any]] = None
+        self.__rendered_str: Mapping[str, str] = {}
+
+    def get_input_map(self) -> Mapping[str, Any]:
+        # This return reference to input map, so input map can be updated
+        return self.__input_map
+
+    def _set_input_map(self, key: str, val: Any):
+        self.__input_map[key] = val
+
+    def get_env_map(self) -> Mapping[str, str]:
+        # This return reference to env map, so env map can be updated
+        return self.__env_map
+
+    def _set_env_map(self, key: str, val: str):
+        self.__env_map[key] = val
+
+    def render_any(
+        self, val: Any, data: Optional[Mapping[str, Any]] = None
+    ) -> Any:
+        if isinstance(val, str):
+            return self.render_str(val, data)
+        return val
+
+    def render_float(
+        self, val: Union[str, float], data: Optional[Mapping[str, Any]] = None
+    ) -> float:
+        if isinstance(val, str):
+            return float(self.render_str(val, data))
+        return val
+
+    def render_int(
+        self, val: Union[str, int], data: Optional[Mapping[str, Any]] = None
+    ) -> int:
+        if isinstance(val, str):
+            return int(self.render_str(val, data))
+        return val
+
+    def render_bool(
+        self, val: Union[str, bool], data: Optional[Mapping[str, Any]] = None
+    ) -> bool:
+        if isinstance(val, str):
+            return to_boolean(self.render_str(val, data))
+        return val
+
+    def render_str(
+        self, val: str, data: Optional[Mapping[str, Any]] = None
+    ) -> str:
+        if val in self.__rendered_str:
+            return self.__rendered_str[val]
+        if not is_probably_jinja(val):
+            return val
+        template = jinja2.Template(val)
+        render_data = self.__get_render_data(additional_data=data)
+        try:
+            rendered_text = template.render(render_data)
+        except Exception:
+            raise Exception(f'Fail to render "{val}" with data: {render_data}')
+        self.__rendered_str[val] = rendered_text
+        return rendered_text
+
+    def render_file(
+        self, location: str, data: Optional[Mapping[str, Any]] = None
+    ) -> str:
+        location_dir = os.path.dirname(location)
+        env = jinja2.Environment(
+            loader=AnyExtensionFileSystemLoader([location_dir])
+        )
+        template = env.get_template(location)
+        render_data = self.__get_render_data(additional_data=data)
+        render_data['TEMPLATE_DIR'] = location_dir
+        rendered_text = template.render(render_data)
+        return rendered_text
+
+    def __get_render_data(
+        self, additional_data: Optional[Mapping[str, Any]] = None
+    ) -> Mapping[str, Any]:
+        self.__ensure_cached_render_data()
+        if additional_data is None:
+            return self.__render_data
+        return {**self.__render_data, **additional_data}
+
+    def __ensure_cached_render_data(self):
+        if self.__render_data is not None:
+            return self.__render_data
+        render_data = dict(DEFAULT_RENDER_DATA)
+        render_data.update({
+            'env': self.__env_map,
+            'input': self.__input_map,
+        })
+        self.__render_data = render_data
+        return render_data