zrb 0.6.0__py3-none-any.whl → 0.6.1__py3-none-any.whl

zrb/task/docker_compose_task.py

@@ -1,56 +1,66 @@
+import os
+import pathlib
+
+from zrb.helper.accessories.name import get_random_name
+from zrb.helper.docker_compose.fetch_external_env import fetch_compose_file_env_map
+from zrb.helper.docker_compose.file import read_compose_file, write_compose_file
+from zrb.helper.string.conversion import to_cli_name
+from zrb.helper.string.modification import double_quote
+from zrb.helper.typecheck import typechecked
 from zrb.helper.typing import (
-    Any, Callable, Iterable, List, Mapping, Optional, Union, TypeVar,
-    JinjaTemplate
+    Any,
+    Callable,
+    Iterable,
+    JinjaTemplate,
+    List,
+    Mapping,
+    Optional,
+    TypeVar,
+    Union,
 )
-from zrb.helper.typecheck import typechecked
-from zrb.task.cmd_task import CmdTask, CmdResult, CmdVal
+from zrb.helper.util import to_snake_case
 from zrb.task.any_task import AnyTask
 from zrb.task.any_task_event_handler import (
-    OnTriggered, OnWaiting, OnSkipped, OnStarted, OnReady, OnRetry, OnFailed
+    OnFailed,
+    OnReady,
+    OnRetry,
+    OnSkipped,
+    OnStarted,
+    OnTriggered,
+    OnWaiting,
 )
+from zrb.task.cmd_task import CmdResult, CmdTask, CmdVal
 from zrb.task_env.constant import RESERVED_ENV_NAMES
 from zrb.task_env.env import Env
 from zrb.task_env.env_file import EnvFile
 from zrb.task_group.group import Group
 from zrb.task_input.any_input import AnyInput
-from zrb.helper.accessories.name import get_random_name
-from zrb.helper.string.conversion import to_cli_name
-from zrb.helper.string.modification import double_quote
-from zrb.helper.docker_compose.file import (
-    read_compose_file, write_compose_file
-)
-from zrb.helper.docker_compose.fetch_external_env import (
-    fetch_compose_file_env_map
-)
-
-import os
-import pathlib
 
 CURRENT_DIR = os.path.dirname(__file__)
-SHELL_SCRIPT_DIR = os.path.join(CURRENT_DIR, '..', 'shell-scripts')
-TDockerComposeTask = TypeVar('TDockerComposeTask', bound='DockerComposeTask')
+SHELL_SCRIPT_DIR = os.path.join(CURRENT_DIR, "..", "shell-scripts")
+TDockerComposeTask = TypeVar("TDockerComposeTask", bound="DockerComposeTask")
 
 ensure_docker_is_installed = CmdTask(
-    name='ensure-docker-is-installed',
+    name="ensure-docker-is-installed",
     cmd_path=[
-        os.path.join(SHELL_SCRIPT_DIR, '_common-util.sh'),
-        os.path.join(SHELL_SCRIPT_DIR, 'ensure-docker-is-installed.sh')
+        os.path.join(SHELL_SCRIPT_DIR, "_common-util.sh"),
+        os.path.join(SHELL_SCRIPT_DIR, "ensure-docker-is-installed.sh"),
     ],
-    preexec_fn=None
+    preexec_fn=None,
 )
 
 ensure_zrb_network_exists = CmdTask(
-    name='ensure-zrb-network-exists',
+    name="ensure-zrb-network-exists",
     cmd=[
-        'docker network inspect zrb >/dev/null 2>&1 || \\',
-        'docker network create -d bridge zrb'
+        "docker network inspect zrb >/dev/null 2>&1 || \\",
+        "docker network create -d bridge zrb",
     ],
     upstreams=[ensure_docker_is_installed],
 )
 
 
 @typechecked
-class ServiceConfig():
+class ServiceConfig:
     def __init__(self, envs: List[Env] = [], env_files: List[EnvFile] = []):
         self._envs = envs
         self._env_files = env_files
@@ -73,17 +83,17 @@ class DockerComposeTask(CmdTask):
         env_files: Iterable[EnvFile] = [],
         icon: Optional[str] = None,
         color: Optional[str] = None,
-        description: str = '',
+        description: str = "",
         executable: Optional[str] = None,
         compose_service_configs: Mapping[str, ServiceConfig] = {},
         compose_file: Optional[str] = None,
-        compose_cmd: str = 'up',
+        compose_cmd: str = "up",
         compose_options: Mapping[JinjaTemplate, JinjaTemplate] = {},
         compose_flags: Iterable[JinjaTemplate] = [],
         compose_args: Iterable[JinjaTemplate] = [],
-        compose_env_prefix: str = '',
-        setup_cmd: CmdVal = '',
-        setup_cmd_path: CmdVal = '',
+        compose_env_prefix: str = "",
+        setup_cmd: CmdVal = "",
+        setup_cmd_path: CmdVal = "",
         cwd: Optional[Union[str, pathlib.Path]] = None,
         upstreams: Iterable[AnyTask] = [],
         on_triggered: Optional[OnTriggered] = None,
@@ -101,7 +111,7 @@ class DockerComposeTask(CmdTask):
         max_error_line: int = 1000,
         preexec_fn: Optional[Callable[[], Any]] = os.setsid,
         should_execute: Union[bool, str, Callable[..., bool]] = True,
-        return_upstream_result: bool = False
+        return_upstream_result: bool = False,
     ):
         combined_env_files = list(env_files)
         CmdTask.__init__(
@@ -132,23 +142,21 @@ class DockerComposeTask(CmdTask):
             max_error_line=max_error_line,
             preexec_fn=preexec_fn,
             should_execute=should_execute,
-            return_upstream_result=return_upstream_result
+            return_upstream_result=return_upstream_result,
         )
         self._setup_cmd = setup_cmd
         self._setup_cmd_path = setup_cmd_path
         self._compose_service_configs = compose_service_configs
         self._compose_cmd = compose_cmd
         self._compose_options = compose_options
-        if '--f' in compose_options:
-            raise ValueError('compose option cannot contains --f')
-        if '--file' in compose_options:
-            raise ValueError('compose option cannot contains --file')
+        if "--f" in compose_options:
+            raise ValueError("compose option cannot contains --f")
+        if "--file" in compose_options:
+            raise ValueError("compose option cannot contains --file")
         self._compose_flags = compose_flags
         self._compose_args = compose_args
         self._compose_env_prefix = compose_env_prefix
-        self._compose_template_file = self.__get_compose_template_file(
-            compose_file
-        )
+        self._compose_template_file = self.__get_compose_template_file(compose_file)
         self._compose_runtime_file = self.__get_compose_runtime_file(
             self._compose_template_file
         )
@@ -184,8 +192,8 @@ class DockerComposeTask(CmdTask):
                 continue
             added_env_map[key] = True
             os_name = key
-            if self._compose_env_prefix != '':
-                os_name = f'{self._compose_env_prefix}_{os_name}'
+            if self._compose_env_prefix != "":
+                os_name = f"{self._compose_env_prefix}_{os_name}"
             self.insert_env(Env(name=key, os_name=os_name, default=value))
 
     def inject_env_files(self):
@@ -201,87 +209,96 @@ class DockerComposeTask(CmdTask):
             for env_file in env_files:
                 envs += env_file.get_envs()
             envs += service_config.get_envs()
-            compose_data = self.__apply_service_env(
-                compose_data, service, envs
-            )
+            compose_data = self.__apply_service_env(compose_data, service, envs)
         write_compose_file(self._compose_runtime_file, compose_data)
 
     def __apply_service_env(
-        self, compose_data: Any, service: str, envs: List[Env]
+        self, compose_data: Any, service_name: str, envs: List[Env]
     ) -> Any:
         # service not found
-        if 'services' not in compose_data or service not in compose_data['services']: # noqa
-            self.log_error(f'Cannot find services.{service}')
+        service_map = compose_data["services"]
+        if "services" not in compose_data or service_name not in service_map:
+            self.log_error(f"Cannot find services.{service_name}")
             return compose_data
         # service has no environment definition
-        if 'environment' not in compose_data['services'][service]:
-            compose_data['services'][service]['environment'] = {
-                env.get_name(): self.__get_env_compose_value(env)
+        if "environment" not in compose_data["services"][service_name]:
+            compose_data["services"][service_name]["environment"] = {
+                env.get_name(): self.__get_env_compose_value(service_name, env)
                 for env in envs
             }
             return compose_data
         # service environment is a map
-        if isinstance(compose_data['services'][service]['environment'], dict):
+        if isinstance(compose_data["services"][service_name]["environment"], dict):
             new_env_map = self.__get_service_new_env_map(
-                compose_data['services'][service]['environment'], envs
+                service_name=service_name,
+                service_env_map=compose_data["services"][service_name]["environment"],
+                new_envs=envs,
             )
             for key, value in new_env_map.items():
-                compose_data['services'][service]['environment'][key] = value
+                compose_data["services"][service_name]["environment"][key] = value
             return compose_data
         # service environment is a list
-        if isinstance(compose_data['services'][service]['environment'], list):
+        if isinstance(compose_data["services"][service_name]["environment"], list):
             new_env_list = self.__get_service_new_env_list(
-                compose_data['services'][service]['environment'], envs
+                service_name=service_name,
+                service_env_list=compose_data["services"][service_name]["environment"],
+                new_envs=envs,
             )
-            compose_data['services'][service]['environment'] += new_env_list
+            compose_data["services"][service_name]["environment"] += new_env_list
             return compose_data
         return compose_data
 
     def __get_service_new_env_map(
-        self, service_env_map: Mapping[str, str], new_envs: List[Env]
+        self, service_name: str, service_env_map: Mapping[str, str], new_envs: List[Env]
     ) -> Mapping[str, str]:
         new_service_envs: Mapping[str, str] = {}
         for env in new_envs:
             env_name = env.get_name()
             if env_name in service_env_map:
                 continue
-            new_service_envs[env_name] = self.__get_env_compose_value(env)
+            new_service_envs[env_name] = self.__get_env_compose_value(service_name, env)
         return new_service_envs
 
     def __get_service_new_env_list(
-        self, service_env_list: List[str], new_envs: List[Env]
+        self, service_name: str, service_env_list: List[str], new_envs: List[Env]
     ) -> List[str]:
         new_service_envs: List[str] = []
         for env in new_envs:
-            should_be_added = 0 == len([
-                service_env for service_env in service_env_list
-                if service_env.startswith(env.get_name() + '=')
-            ])
+            should_be_added = 0 == len(
+                [
+                    service_env
+                    for service_env in service_env_list
+                    if service_env.startswith(env.get_name() + "=")
+                ]
+            )
             if not should_be_added:
                 continue
             new_service_envs.append(
-                env.get_name() + '=' + self.__get_env_compose_value(env)
+                env.get_name() + "=" + self.__get_env_compose_value(service_name, env)
             )
         return new_service_envs
 
-    def __get_env_compose_value(self, env: Env) -> str:
-        return '${' + env.get_name() + ':-' + env.get_default() + '}'
+    def __get_env_compose_value(self, service_name: str, env: Env) -> str:
+        env_prefix = to_snake_case(service_name).upper()
+        env_name = env.get_name()
+        env_default = env.get_default()
+        return "".join(["${", f"{env_prefix}_{env_name}:-{env_default}", "}"])
 
     def __get_compose_runtime_file(self, compose_file_name: str) -> str:
         directory, file = os.path.split(compose_file_name)
-        prefix = '_' if file.startswith('.') else '._'
+        prefix = "_" if file.startswith(".") else "._"
         runtime_prefix = self.get_cli_name()
         if self._group is not None:
             group_prefix = to_cli_name(self._group._get_full_cli_name())
-            runtime_prefix = f'{group_prefix}-{runtime_prefix}'
-        runtime_prefix += '-' + get_random_name(
-            separator='-', add_random_digit=True, digit_count=3
+            runtime_prefix = f"{group_prefix}-{runtime_prefix}"
+        runtime_prefix += "-" + get_random_name(
+            separator="-", add_random_digit=True, digit_count=3
         )
-        runtime_prefix = '.' + runtime_prefix + '.runtime'
-        file_parts = file.split('.')
+        runtime_prefix = "." + runtime_prefix + ".runtime"
+        file_parts = file.split(".")
         if len(file_parts) > 1:
             file_parts[-2] += runtime_prefix
-            runtime_file_name = prefix + '.'.join(file_parts)
+            runtime_file_name = prefix + ".".join(file_parts)
             return os.path.join(directory, runtime_file_name)
         runtime_file_name = prefix + file + runtime_prefix
         return os.path.join(directory, runtime_file_name)
@@ -289,40 +306,52 @@ class DockerComposeTask(CmdTask):
     def __get_compose_template_file(self, compose_file: Optional[str]) -> str:
         if compose_file is None:
             for _compose_file in [
-                'compose.yml', 'compose.yaml',
-                'docker-compose.yml', 'docker-compose.yaml'
+                "compose.yml",
+                "compose.yaml",
+                "docker-compose.yml",
+                "docker-compose.yaml",
             ]:
                 if os.path.exists(os.path.join(self._cwd, _compose_file)):
                     return os.path.join(self._cwd, _compose_file)
-            raise Exception(f'Cannot find compose file on {self._cwd}')
+            raise Exception(f"Cannot find compose file on {self._cwd}")
         if os.path.isabs(compose_file) and os.path.exists(compose_file):
             return compose_file
         if os.path.exists(os.path.join(self._cwd, compose_file)):
             return os.path.join(self._cwd, compose_file)
-        raise Exception(f'Invalid compose file: {compose_file}')
+        raise Exception(f"Invalid compose file: {compose_file}")
 
     def get_cmd_script(self, *args: Any, **kwargs: Any) -> str:
         setup_cmd_str = self._create_cmd_script(
             self._setup_cmd_path, self._setup_cmd, *args, **kwargs
         )
         command_options = dict(self._compose_options)
-        command_options['--file'] = self._compose_runtime_file
-        options = ' '.join([
-            f'{self.render_str(key)} {double_quote(self.render_str(val))}'
-            for key, val in command_options.items()
-            if self.render_str(val) != ''
-        ])
-        flags = ' '.join([
-            self.render_str(flag) for flag in self._compose_flags
-            if self.render_str(flag) != ''
-        ])
-        args = ' '.join([
-            double_quote(self.render_str(arg)) for arg in self._compose_args
-            if self.render_str(arg) != ''
-        ])
-        cmd_str = '\n'.join([
-            setup_cmd_str,
-            f'docker compose {options} {self._compose_cmd} {flags} {args}',
-        ])
-        self.log_info(f'Command: {cmd_str}')
+        command_options["--file"] = self._compose_runtime_file
+        options = " ".join(
+            [
+                f"{self.render_str(key)} {double_quote(self.render_str(val))}"
+                for key, val in command_options.items()
+                if self.render_str(val) != ""
+            ]
+        )
+        flags = " ".join(
+            [
+                self.render_str(flag)
+                for flag in self._compose_flags
+                if self.render_str(flag) != ""
+            ]
+        )
+        args = " ".join(
+            [
+                double_quote(self.render_str(arg))
+                for arg in self._compose_args
+                if self.render_str(arg) != ""
+            ]
+        )
+        cmd_str = "\n".join(
+            [
+                setup_cmd_str,
+                f"docker compose {options} {self._compose_cmd} {flags} {args}",
+            ]
+        )
+        self.log_info(f"Command: {cmd_str}")
         return cmd_str

zrb/task/flow_task.py

@@ -1,24 +1,26 @@
-from zrb.helper.typing import (
-    Callable, Iterable, List, Optional, TypeVar, Union
-)
 from zrb.helper.typecheck import typechecked
-from zrb.task.base_task.base_task import BaseTask
+from zrb.helper.typing import Callable, Iterable, List, Optional, TypeVar, Union
 from zrb.task.any_task import AnyTask
 from zrb.task.any_task_event_handler import (
-    OnTriggered, OnWaiting, OnSkipped, OnStarted, OnReady, OnRetry, OnFailed
+    OnFailed,
+    OnReady,
+    OnRetry,
+    OnSkipped,
+    OnStarted,
+    OnTriggered,
+    OnWaiting,
 )
+from zrb.task.base_task.base_task import BaseTask
 from zrb.task_env.env import Env
 from zrb.task_env.env_file import EnvFile
 from zrb.task_group.group import Group
 from zrb.task_input.any_input import AnyInput
 
-
-TFlowTask = TypeVar('TFlowTask', bound='FlowTask')
+TFlowTask = TypeVar("TFlowTask", bound="FlowTask")
 
 
 @typechecked
 class FlowTask(BaseTask):
-
     def __init__(
         self,
         name: str,
@@ -28,7 +30,7 @@ class FlowTask(BaseTask):
         env_files: Iterable[EnvFile] = [],
         icon: Optional[str] = None,
         color: Optional[str] = None,
-        description: str = '',
+        description: str = "",
         upstreams: Iterable[AnyTask] = [],
         on_triggered: Optional[OnTriggered] = None,
         on_waiting: Optional[OnWaiting] = None,
@@ -43,7 +45,7 @@ class FlowTask(BaseTask):
         retry_interval: float = 1,
         steps: List[Union[AnyTask, List[AnyTask]]] = [],
         should_execute: Union[bool, str, Callable[..., bool]] = True,
-        return_upstream_result: bool = False
+        return_upstream_result: bool = False,
     ):
         final_upstreams: List[AnyTask] = list(upstreams)
         inputs: List[AnyInput] = list(inputs)
@@ -56,7 +58,7 @@ class FlowTask(BaseTask):
                 upstreams=final_upstreams,
                 inputs=inputs,
                 envs=envs,
-                env_files=env_files
+                env_files=env_files,
             )
             final_upstreams = new_upstreams
         BaseTask.__init__(
@@ -83,15 +85,13 @@ class FlowTask(BaseTask):
             retry_interval=retry_interval,
             should_execute=should_execute,
             return_upstream_result=return_upstream_result,
-            run=lambda *args, **kwargs: kwargs.get('_task').print_out('🆗')
+            run=lambda *args, **kwargs: kwargs.get("_task").print_out("🆗"),
         )
 
     def copy(self) -> TFlowTask:
         return super().copy()
 
-    def _step_to_tasks(
-        self, node: Union[AnyTask, List[AnyTask]]
-    ) -> List[AnyTask]:
+    def _step_to_tasks(self, node: Union[AnyTask, List[AnyTask]]) -> List[AnyTask]:
         if isinstance(node, AnyTask):
             return [node]
         return node
@@ -107,9 +107,7 @@ class FlowTask(BaseTask):
         embeded_tasks: List[AnyTask] = []
         for task in tasks:
             embeded_task = task.copy()
-            embeded_task_root_upstreams = self._get_root_upstreams(
-                tasks=[embeded_task]
-            )
+            embeded_task_root_upstreams = self._get_root_upstreams(tasks=[embeded_task])
             for embeded_task_root_upstream in embeded_task_root_upstreams:
                 embeded_task_root_upstream.add_upstream(*upstreams)
             # embeded_task.add_upstream(*upstreams)
@@ -132,4 +130,3 @@ class FlowTask(BaseTask):
                     continue
                 root_upstreams += self._get_root_upstreams([upstream])
         return root_upstreams
-

zrb/task/http_checker.py

@@ -1,23 +1,36 @@
+from http.client import HTTPConnection, HTTPSConnection
+
+from zrb.helper.typecheck import typechecked
 from zrb.helper.typing import (
-    Any, Callable, Iterable, Optional, Union, TypeVar, JinjaTemplate
+    Any,
+    Callable,
+    Iterable,
+    JinjaTemplate,
+    Optional,
+    TypeVar,
+    Union,
 )
-from zrb.helper.typecheck import typechecked
-from zrb.task.checker import Checker
-from http.client import HTTPConnection, HTTPSConnection
 from zrb.task.any_task import AnyTask
 from zrb.task.any_task_event_handler import (
-    OnTriggered, OnWaiting, OnSkipped, OnStarted, OnReady, OnRetry, OnFailed
+    OnFailed,
+    OnReady,
+    OnRetry,
+    OnSkipped,
+    OnStarted,
+    OnTriggered,
+    OnWaiting,
 )
+from zrb.task.checker import Checker
 from zrb.task_env.env import Env
 from zrb.task_env.env_file import EnvFile
 from zrb.task_group.group import Group
 from zrb.task_input.any_input import AnyInput
 
-THTTPChecker = TypeVar('THTTPChecker', bound='HTTPChecker')
+THTTPChecker = TypeVar("THTTPChecker", bound="HTTPChecker")
 
 
 @typechecked
-class HttpConnectionConfig():
+class HttpConnectionConfig:
     def __init__(
         self,
         is_https: bool,
@@ -31,40 +44,35 @@ class HttpConnectionConfig():
         self.method = method
         self.host = host
         self.port = port
-        self.url = '/' + url if not url.startswith('/') else url
+        self.url = "/" + url if not url.startswith("/") else url
         self.timeout = timeout
-        self.protocol = 'https' if self.is_https else 'http'
-        full_url = f'{self.protocol}://{self.host}:{self.port}{self.url}'
-        self.label = f'{self.method} {full_url}'
+        self.protocol = "https" if self.is_https else "http"
+        full_url = f"{self.protocol}://{self.host}:{self.port}{self.url}"
+        self.label = f"{self.method} {full_url}"
 
     def get_connection(self) -> Union[HTTPConnection, HTTPSConnection]:
         if self.is_https:
-            return HTTPSConnection(
-                host=self.host, port=self.port, timeout=self.timeout
-            )
-        return HTTPConnection(
-            host=self.host, port=self.port, timeout=self.timeout
-        )
+            return HTTPSConnection(host=self.host, port=self.port, timeout=self.timeout)
+        return HTTPConnection(host=self.host, port=self.port, timeout=self.timeout)
 
 
 @typechecked
 class HTTPChecker(Checker):
-
     def __init__(
         self,
-        name: str = 'check-http',
+        name: str = "check-http",
         group: Optional[Group] = None,
         inputs: Iterable[AnyInput] = [],
         envs: Iterable[Env] = [],
         env_files: Iterable[EnvFile] = [],
         icon: Optional[str] = None,
         color: Optional[str] = None,
-        description: str = '',
-        host: JinjaTemplate = 'localhost',
+        description: str = "",
+        host: JinjaTemplate = "localhost",
         port: Union[int, JinjaTemplate] = 80,
         timeout: Union[int, JinjaTemplate] = 5,
-        method: str = 'HEAD',
-        url: str = '/',
+        method: str = "HEAD",
+        url: str = "/",
         is_https: Union[bool, str] = False,
         upstreams: Iterable[AnyTask] = [],
         on_triggered: Optional[OnTriggered] = None,
@@ -77,7 +85,7 @@ class HTTPChecker(Checker):
         checking_interval: Union[int, float] = 0.1,
         progress_interval: Union[int, float] = 5,
         expected_result: bool = True,
-        should_execute: Union[bool, JinjaTemplate, Callable[..., bool]] = True
+        should_execute: Union[bool, JinjaTemplate, Callable[..., bool]] = True,
     ):
         Checker.__init__(
             self,
@@ -115,14 +123,12 @@ class HTTPChecker(Checker):
 
     def to_function(
         self,
-        env_prefix: str = '',
+        env_prefix: str = "",
         raise_error: bool = True,
         is_async: bool = False,
-        show_done_info: bool = True
+        show_done_info: bool = True,
     ) -> Callable[..., bool]:
-        return super().to_function(
-            env_prefix, raise_error, is_async, show_done_info
-        )
+        return super().to_function(env_prefix, raise_error, is_async, show_done_info)
 
     async def run(self, *args: Any, **kwargs: Any) -> bool:
         self._config = HttpConnectionConfig(
@@ -131,7 +137,7 @@ class HTTPChecker(Checker):
             host=self.render_str(self._host),
             port=self.render_int(self._port),
             url=self.render_str(self._url),
-            timeout=self.render_int(self._timeout)
+            timeout=self.render_int(self._timeout),
         )
         return await super().run(*args, **kwargs)
 
@@ -141,14 +147,12 @@ class HTTPChecker(Checker):
             conn.request(self._config.method, self._config.url)
             res = conn.getresponse()
             if res.status < 300:
-                self.log_info('Connection success')
-                self.print_out(f'{self._config.label} {res.status} (OK)')
+                self.log_info("Connection success")
+                self.print_out(f"{self._config.label} {res.status} (OK)")
                 return True
-            self.show_progress(
-                f'{self._config.label} {res.status} (Not OK)'
-            )
+            self.show_progress(f"{self._config.label} {res.status} (Not OK)")
         except Exception:
-            self.show_progress(f'{self._config.label} Connection error')
+            self.show_progress(f"{self._config.label} Connection error")
         finally:
             conn.close()
         return False