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

zrb/task/checker.py

@@ -1,31 +1,36 @@
-from zrb.helper.typing import Any, Callable, Iterable, Optional, Union
+import asyncio
+
 from zrb.helper.typecheck import typechecked
-from zrb.task.base_task.base_task import BaseTask
+from zrb.helper.typing import Any, Callable, Iterable, Optional, 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
 
-import asyncio
-
 
 @typechecked
 class Checker(BaseTask):
-
     def __init__(
         self,
-        name: str = 'check',
+        name: str = "check",
         group: Optional[Group] = None,
         inputs: Iterable[AnyInput] = [],
         envs: Iterable[Env] = [],
         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,
@@ -37,7 +42,7 @@ class Checker(BaseTask):
         checking_interval: Union[int, float] = 0.1,
         progress_interval: Union[int, float] = 30,
         expected_result: bool = True,
-        should_execute: Union[bool, str, Callable[..., bool]] = True
+        should_execute: Union[bool, str, Callable[..., bool]] = True,
     ):
         BaseTask.__init__(
             self,

zrb/task/cmd_task.py

@@ -1,31 +1,44 @@
+import asyncio
+import atexit
+import os
+import pathlib
+import signal
+import subprocess
+import sys
+import time
+
+from zrb.config.config import default_shell
+from zrb.helper.string.conversion import to_variable_name
+from zrb.helper.typecheck import typechecked
 from zrb.helper.typing import (
-    Any, Callable, Iterable, List, Optional, Union, TypeVar, JinjaTemplate
+    Any,
+    Callable,
+    Iterable,
+    JinjaTemplate,
+    List,
+    Optional,
+    TypeVar,
+    Union,
 )
-from zrb.helper.typecheck import typechecked
-from zrb.helper.string.conversion import to_variable_name
 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
-from zrb.config.config import default_shell
-
-import asyncio
-import atexit
-import os
-import pathlib
-import signal
-import sys
-import subprocess
-import time
 
 _has_stty = True
 try:
-    _original_stty = subprocess.getoutput('stty -g').rstrip()
+    _original_stty = subprocess.getoutput("stty -g").rstrip()
 except Exception:
     _has_stty = False
 
@@ -34,7 +47,7 @@ def _reset_stty():
     global _has_stty
     if _has_stty:
         try:
-            subprocess.run(['stty', _original_stty], check=True)
+            subprocess.run(["stty", _original_stty], check=True)
         except Exception:
             _has_stty = False
 
@@ -42,12 +55,12 @@ def _reset_stty():
 CmdVal = Union[
     JinjaTemplate,
     Iterable[JinjaTemplate],
-    Callable[..., Union[Iterable[JinjaTemplate], JinjaTemplate]]
+    Callable[..., Union[Iterable[JinjaTemplate], JinjaTemplate]],
 ]
-TCmdTask = TypeVar('TCmdTask', bound='CmdTask')
+TCmdTask = TypeVar("TCmdTask", bound="CmdTask")
 
 
-class CmdResult():
+class CmdResult:
     def __init__(self, output: str, error: str):
         self.output = output
         self.error = error
@@ -56,7 +69,7 @@ class CmdResult():
         return self.output
 
 
-class CmdGlobalState():
+class CmdGlobalState:
     def __init__(self):
         self.no_more_attempt: bool = False
         self.is_killed_by_signal: bool = False
@@ -64,7 +77,7 @@ class CmdGlobalState():
 
 @typechecked
 class CmdTask(BaseTask):
-    '''
+    """
     Command Task.
     You can use this task to run shell command.
 
@@ -80,7 +93,7 @@ class CmdTask(BaseTask):
         >>>     ]
         >>> )
         >>> runner.register(hello)
-    '''
+    """
 
     _pids: List[int] = []
     _global_state = CmdGlobalState()
@@ -94,10 +107,10 @@ class CmdTask(BaseTask):
         env_files: Iterable[EnvFile] = [],
         icon: Optional[str] = None,
         color: Optional[str] = None,
-        description: str = '',
+        description: str = "",
         executable: Optional[str] = None,
-        cmd: CmdVal = '',
-        cmd_path: CmdVal = '',
+        cmd: CmdVal = "",
+        cmd_path: CmdVal = "",
         cwd: Optional[Union[str, pathlib.Path]] = None,
         upstreams: Iterable[AnyTask] = [],
         on_triggered: Optional[OnTriggered] = None,
@@ -115,7 +128,7 @@ class CmdTask(BaseTask):
         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,
     ):
         BaseTask.__init__(
             self,
@@ -140,7 +153,7 @@ class CmdTask(BaseTask):
             retry=retry,
             retry_interval=retry_interval,
             should_execute=should_execute,
-            return_upstream_result=return_upstream_result
+            return_upstream_result=return_upstream_result,
         )
         max_output_line = max_output_line if max_output_line > 0 else 1
         max_error_line = max_error_line if max_error_line > 0 else 1
@@ -151,7 +164,7 @@ class CmdTask(BaseTask):
         self._max_error_size = max_error_line
         self._output_buffer: Iterable[str] = []
         self._error_buffer: Iterable[str] = []
-        if executable is None and default_shell != '':
+        if executable is None and default_shell != "":
             executable = default_shell
         self._executable = executable
         self._process: Optional[asyncio.subprocess.Process]
@@ -163,9 +176,7 @@ class CmdTask(BaseTask):
     def set_cwd(self, cwd: Union[str, pathlib.Path]):
         self.__set_cwd(cwd)
 
-    def __set_cwd(
-        self, cwd: Optional[Union[str, pathlib.Path]]
-    ):
+    def __set_cwd(self, cwd: Optional[Union[str, pathlib.Path]]):
         if cwd is None:
             self._cwd: Union[str, pathlib.Path] = os.getcwd()
             return
@@ -173,17 +184,15 @@ class CmdTask(BaseTask):
 
     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[..., CmdResult]:
-        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)
 
     def print_result(self, result: CmdResult):
-        if result.output == '':
+        if result.output == "":
             return
         print(result.output)
 
@@ -193,19 +202,21 @@ class CmdTask(BaseTask):
         for task_input in self._get_combined_inputs():
             input_key = to_variable_name(task_input.get_name())
             input_value = input_map.get(input_key)
-            env_name = '_INPUT_' + input_key.upper()
+            env_name = "_INPUT_" + input_key.upper()
             should_render = task_input.should_render()
-            self.add_env(Env(
-                name=env_name,
-                os_name='',
-                default=str(input_value),
-                should_render=should_render
-            ))
+            self.add_env(
+                Env(
+                    name=env_name,
+                    os_name="",
+                    default=str(input_value),
+                    should_render=should_render,
+                )
+            )
 
     async def run(self, *args: Any, **kwargs: Any) -> CmdResult:
         cmd = self.get_cmd_script(*args, **kwargs)
-        self.print_out_dark('Run script: ' + self.__get_multiline_repr(cmd))
-        self.print_out_dark('Working directory: ' + self._cwd)
+        self.print_out_dark("Run script: " + self.__get_multiline_repr(cmd))
+        self.print_out_dark("Working directory: " + self._cwd)
         self._output_buffer = []
         self._error_buffer = []
         process = await asyncio.create_subprocess_shell(
@@ -218,7 +229,7 @@ class CmdTask(BaseTask):
             executable=self._executable,
             close_fds=True,
             preexec_fn=self._preexec_fn,
-            bufsize=0
+            bufsize=0,
         )
         self._set_task_pid(process.pid)
         self._pids.append(process.pid)
@@ -230,19 +241,17 @@ class CmdTask(BaseTask):
             self.print_err(e)
         atexit.register(self.__on_exit)
         await self.__wait_process(process)
-        self.log_info('Process completed')
+        self.log_info("Process completed")
         atexit.unregister(self.__on_exit)
-        output = '\n'.join(self._output_buffer)
-        error = '\n'.join(self._error_buffer)
+        output = "\n".join(self._output_buffer)
+        error = "\n".join(self._error_buffer)
         # get return code
         return_code = process.returncode
-        self.log_info(f'Exit status: {return_code}')
+        self.log_info(f"Exit status: {return_code}")
         if return_code != 0 and not self._global_state.is_killed_by_signal:
-            raise Exception(
-                f'Process {self._name} exited ({return_code}): {error}'
-            )
-        self.set_task_xcom(key='output', value=output)
-        self.set_task_xcom(key='error', value=error)
+            raise Exception(f"Process {self._name} exited ({return_code}): {error}")
+        self.set_task_xcom(key="output", value=output)
+        self.set_task_xcom(key="error", value=error)
         return CmdResult(output, error)
 
     def _should_attempt(self) -> bool:
@@ -258,7 +267,7 @@ class CmdTask(BaseTask):
     def __on_kill(self, signum: Any, frame: Any):
         self._global_state.no_more_attempt = True
         self._global_state.is_killed_by_signal = True
-        self.print_out_dark(f'Getting signal {signum}')
+        self.print_out_dark(f"Getting signal {signum}")
         for pid in self._pids:
             self.__kill_by_pid(pid)
         tasks = asyncio.all_tasks()
@@ -268,7 +277,7 @@ class CmdTask(BaseTask):
             except Exception as e:
                 self.print_err(e)
         time.sleep(0.3)
-        self.print_out_dark(f'Exiting with signal {signum}')
+        self.print_out_dark(f"Exiting with signal {signum}")
         sys.exit(signum)
 
     def __on_exit(self):
@@ -276,27 +285,27 @@ class CmdTask(BaseTask):
         self.__kill_by_pid(self._process.pid)
 
     def __kill_by_pid(self, pid: int):
-        '''
+        """
         Kill a pid, gracefully
-        '''
+        """
         try:
             process_ever_exists = False
             if self.__is_process_exist(pid):
                 process_ever_exists = True
-                self.print_out_dark(f'Send SIGTERM to process {pid}')
+                self.print_out_dark(f"Send SIGTERM to process {pid}")
                 os.killpg(os.getpgid(pid), signal.SIGTERM)
                 time.sleep(0.3)
             if self.__is_process_exist(pid):
-                self.print_out_dark(f'Send SIGINT to process {pid}')
+                self.print_out_dark(f"Send SIGINT to process {pid}")
                 os.killpg(os.getpgid(pid), signal.SIGINT)
                 time.sleep(0.3)
             if self.__is_process_exist(pid):
-                self.print_out_dark(f'Send SIGKILL to process {pid}')
+                self.print_out_dark(f"Send SIGKILL to process {pid}")
                 os.killpg(os.getpgid(pid), signal.SIGKILL)
             if process_ever_exists:
-                self.print_out_dark(f'Process {pid} is killed successfully')
+                self.print_out_dark(f"Process {pid} is killed successfully")
         except Exception:
-            self.log_error(f'Cannot kill process {pid}')
+            self.log_error(f"Cannot kill process {pid}")
 
     def __is_process_exist(self, pid: int) -> bool:
         try:
@@ -310,21 +319,23 @@ class CmdTask(BaseTask):
         stdout_queue = asyncio.Queue()
         stderr_queue = asyncio.Queue()
         # Read from streams and put into queue
-        stdout_process = asyncio.create_task(self.__queue_stream(
-            process.stdout, stdout_queue
-        ))
-        stderr_process = asyncio.create_task(self.__queue_stream(
-            process.stderr, stderr_queue
-        ))
+        stdout_process = asyncio.create_task(
+            self.__queue_stream(process.stdout, stdout_queue)
+        )
+        stderr_process = asyncio.create_task(
+            self.__queue_stream(process.stderr, stderr_queue)
+        )
         # Handle messages in queue
-        stdout_log_process = asyncio.create_task(self.__log_from_queue(
-            stdout_queue, self.print_out,
-            self._output_buffer, self._max_output_size
-        ))
-        stderr_log_process = asyncio.create_task(self.__log_from_queue(
-            stderr_queue, self.print_err,
-            self._error_buffer, self._max_error_size
-        ))
+        stdout_log_process = asyncio.create_task(
+            self.__log_from_queue(
+                stdout_queue, self.print_out, self._output_buffer, self._max_output_size
+            )
+        )
+        stderr_log_process = asyncio.create_task(
+            self.__log_from_queue(
+                stderr_queue, self.print_err, self._error_buffer, self._max_error_size
+            )
+        )
         # wait process
         await process.wait()
         # wait reader and logger
@@ -336,14 +347,12 @@ class CmdTask(BaseTask):
         await stderr_log_process
 
     def get_cmd_script(self, *args: Any, **kwargs: Any) -> str:
-        return self._create_cmd_script(
-            self._cmd_path, self._cmd, *args, **kwargs
-        )
+        return self._create_cmd_script(self._cmd_path, self._cmd, *args, **kwargs)
 
     def _create_cmd_script(
         self, cmd_path: CmdVal, cmd: CmdVal, *args: Any, **kwargs: Any
     ) -> str:
-        if not isinstance(cmd_path, str) or cmd_path != '':
+        if not isinstance(cmd_path, str) or cmd_path != "":
             if callable(cmd_path):
                 return self.__get_rendered_cmd_path(cmd_path(*args, **kwargs))
             return self.__get_rendered_cmd_path(cmd_path)
@@ -351,22 +360,17 @@ class CmdTask(BaseTask):
             return self.__get_rendered_cmd(cmd(*args, **kwargs))
         return self.__get_rendered_cmd(cmd)
 
-    def __get_rendered_cmd_path(
-        self, cmd_path: Union[str, Iterable[str]]
-    ) -> str:
+    def __get_rendered_cmd_path(self, cmd_path: Union[str, Iterable[str]]) -> str:
         if isinstance(cmd_path, str):
             return self.render_file(cmd_path)
-        return '\n'.join([
-            self.render_file(cmd_path_str)
-            for cmd_path_str in cmd_path
-        ])
+        return "\n".join([self.render_file(cmd_path_str) for cmd_path_str in cmd_path])
 
     def __get_rendered_cmd(
         self, cmd: Union[JinjaTemplate, Iterable[JinjaTemplate]]
     ) -> str:
         if isinstance(cmd, str):
             return self.render_str(cmd)
-        return self.render_str('\n'.join(list(cmd)))
+        return self.render_str("\n".join(list(cmd)))
 
     async def __queue_stream(self, stream, queue: asyncio.Queue):
         while True:
@@ -383,31 +387,29 @@ class CmdTask(BaseTask):
         queue: asyncio.Queue,
         print_log: Callable[[str], None],
         buffer: Iterable[str],
-        max_size: int
+        max_size: int,
     ):
         while True:
             line = await queue.get()
             if not line:
                 break
-            line_str = line.decode('utf-8').rstrip()
+            line_str = line.decode("utf-8").rstrip()
             self.__add_to_buffer(buffer, max_size, line_str)
             _reset_stty()
             print_log(line_str)
             _reset_stty()
 
-    def __add_to_buffer(
-        self, buffer: Iterable[str], max_size: int, new_line: str
-    ):
+    def __add_to_buffer(self, buffer: Iterable[str], max_size: int, new_line: str):
         if len(buffer) >= max_size:
             buffer.pop(0)
         buffer.append(new_line)
 
     def __get_multiline_repr(self, text: str) -> str:
         lines_repr: Iterable[str] = []
-        lines = text.split('\n')
+        lines = text.split("\n")
         if len(lines) == 1:
             return lines[0]
         for index, line in enumerate(lines):
-            line_number_repr = str(index + 1).rjust(4, '0')
-            lines_repr.append(f'        {line_number_repr} | {line}')
-        return '\n' + '\n'.join(lines_repr)
+            line_number_repr = str(index + 1).rjust(4, "0")
+            lines_repr.append(f"        {line_number_repr} | {line}")
+        return "\n" + "\n".join(lines_repr)

zrb/task/decorator.py

@@ -1,17 +1,21 @@
-from zrb.helper.typing import (
-    Any, Callable, Iterable, Optional, Union
-)
+from zrb.action.runner import Runner
 from zrb.helper.typecheck import typechecked
+from zrb.helper.typing import Any, Callable, Iterable, Optional, 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_input.any_input import AnyInput
+from zrb.task.task import Task
 from zrb.task_env.env import Env
 from zrb.task_env.env_file import EnvFile
 from zrb.task_group.group import Group
-from zrb.action.runner import Runner
-from zrb.task.task import Task
+from zrb.task_input.any_input import AnyInput
 
 # flake8: noqa E501
 
@@ -25,7 +29,7 @@ def python_task(
     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,
@@ -40,14 +44,14 @@ def python_task(
     retry_interval: float = 1,
     should_execute: Union[bool, str, Callable[..., bool]] = True,
     return_upstream_result: bool = False,
-    runner: Optional[Runner] = None
+    runner: Optional[Runner] = None,
 ) -> Callable[[Callable[..., Any]], Task]:
-    '''
+    """
     python_task decorator helps you turn any Python function into a task
 
     Returns:
         Callable[[Callable[..., Any]], Task]: A callable turning function into task.
-    
+
     Examples:
         >>> from zrb import python_task
         >>> @python_task(
@@ -57,7 +61,8 @@ def python_task(
         >>>    return 'hello world'
         >>> print(my_task)
         <Task name=my-task>
-    '''
+    """
+
     def _create_task(fn: Callable[..., Any]) -> Task:
         task = Task(
             name=name,
@@ -82,9 +87,10 @@ def python_task(
             retry_interval=retry_interval,
             run=fn,
             should_execute=should_execute,
-            return_upstream_result=return_upstream_result
+            return_upstream_result=return_upstream_result,
         )
         if runner is not None:
             runner.register(task)
         return task
+
     return _create_task