zrb 1.0.0b8__py3-none-any.whl → 1.0.0b10__py3-none-any.whl

zrb/builtin/project/add/fastapp/fastapp_template/my_app_name/test/test_not_found_error.py

@@ -0,0 +1,16 @@
+import os
+
+from fastapi.testclient import TestClient
+from my_app_name.main import app
+from my_app_name.module.gateway.util.view import render_error
+
+
+def test_not_found_error():
+    client = TestClient(app, base_url="http://localhost")
+    response = client.get(
+        "/holly-grail/not-found/philosopher-stone/not-found/hope/inexist"
+    )
+    assert response.status_code == 404
+    assert response.text == render_error(
+        error_message="Not found", status_code=404
+    ).body.decode("utf-8")

zrb/builtin/project/add/fastapp/fastapp_template/my_app_name/test.sh

@@ -0,0 +1,7 @@
+pytest -vv \
+ --cov=my_app_name \
+ --cov-config=.coveragerc \
+ --cov-report=html \
+ --cov-report=term \
+ --cov-report=term-missing \
+ --ignore=_zrb

zrb/task/base_task.py

@@ -295,16 +295,16 @@ class BaseTask(AnyTask):
     async def exec_root_tasks(self, session: AnySession):
         session.set_main_task(self)
         session.state_logger.write(session.as_state_log())
-        log_state = asyncio.create_task(self._log_session_state(session))
-        root_tasks = [
-            task
-            for task in session.get_root_tasks(self)
-            if session.is_allowed_to_run(task)
-        ]
-        root_task_coros = [
-            run_async(root_task.exec_chain(session)) for root_task in root_tasks
-        ]
         try:
+            log_state = asyncio.create_task(self._log_session_state(session))
+            root_tasks = [
+                task
+                for task in session.get_root_tasks(self)
+                if session.is_allowed_to_run(task)
+            ]
+            root_task_coros = [
+                run_async(root_task.exec_chain(session)) for root_task in root_tasks
+            ]
             await asyncio.gather(*root_task_coros)
             await session.wait_deferred()
             session.terminate()
@@ -312,7 +312,7 @@ class BaseTask(AnyTask):
             return session.final_result
         except IndexError:
             return None
-        except asyncio.CancelledError:
+        except (asyncio.CancelledError, KeyboardInterrupt):
             ctx = self.get_ctx(session)
             ctx.log_info("Session terminated")
         finally:

zrb/task/cmd_task.py

@@ -141,11 +141,8 @@ class CmdTask(BaseTask):
             max_error_line=self._max_error_line,
         )
         # Check for errors
-        if return_code != 0:
-            ctx.log_error(f"Exit status: {return_code}")
-            raise Exception(
-                f"Process {self._name} exited ({return_code}): {cmd_result.error}"
-            )
+        if return_code > 0:
+            raise Exception(f"Process {self._name} exited ({return_code})")
         ctx.log_info(f"Exit status: {return_code}")
         return cmd_result
 

zrb/util/cmd/command.py

@@ -1,26 +1,12 @@
 import asyncio
 import os
 import re
-import signal
 import sys
 from collections.abc import Callable
 
-from zrb.cmd.cmd_result import CmdResult
-
-_RUNNING_PROCESSES = []
-
-
-def _cleanup_processes():
-    """Terminate all running subprocesses on exit."""
-    for process in _RUNNING_PROCESSES:
-        if process.returncode is None:
-            process.terminate()
-            process.kill()
-
+import psutil
 
-# Register a single cleanup function once
-signal.signal(signal.SIGINT, lambda sig, frame: _cleanup_processes())
-signal.signal(signal.SIGTERM, lambda sig, frame: _cleanup_processes())
+from zrb.cmd.cmd_result import CmdResult
 
 
 def check_unrecommended_commands(cmd_script: str) -> dict[str, str]:
@@ -81,35 +67,40 @@ async def run_command(
 
     actual_print_method = print_method if print_method is not None else print
     cmd_process = None
-    try:
-        if cwd is None:
-            cwd = os.getcwd()
-        if env_map is None:
-            env_map = os.environ
-        cmd_process = await asyncio.create_subprocess_exec(
-            *cmd,
-            cwd=cwd,
-            stdin=sys.stdin if sys.stdin.isatty() else None,
-            stdout=asyncio.subprocess.PIPE,
-            stderr=asyncio.subprocess.PIPE,
-            env=env_map,
-            bufsize=0,
-        )
-        if register_pid_method is not None:
-            register_pid_method(cmd_process.pid)
-        stdout_task = asyncio.create_task(
-            __read_stream(cmd_process.stdout, actual_print_method, max_output_line)
-        )
-        stderr_task = asyncio.create_task(
-            __read_stream(cmd_process.stderr, actual_print_method, max_error_line)
-        )
-        # Wait for process to complete and gather stdout/stderr
-        return_code = await cmd_process.wait()
-        stdout = await stdout_task
-        stderr = await stderr_task
-        return CmdResult(stdout, stderr), return_code
-    finally:
-        if cmd_process in _RUNNING_PROCESSES:
-            _RUNNING_PROCESSES.remove(cmd_process)
-        if cmd_process is not None and cmd_process.returncode is None:
-            cmd_process.terminate()
+    if cwd is None:
+        cwd = os.getcwd()
+    if env_map is None:
+        env_map = os.environ
+    cmd_process = await asyncio.create_subprocess_exec(
+        *cmd,
+        cwd=cwd,
+        stdin=sys.stdin if sys.stdin.isatty() else None,
+        stdout=asyncio.subprocess.PIPE,
+        stderr=asyncio.subprocess.PIPE,
+        env=env_map,
+        bufsize=0,
+    )
+    if register_pid_method is not None:
+        register_pid_method(cmd_process.pid)
+    stdout_task = asyncio.create_task(
+        __read_stream(cmd_process.stdout, actual_print_method, max_output_line)
+    )
+    stderr_task = asyncio.create_task(
+        __read_stream(cmd_process.stderr, actual_print_method, max_error_line)
+    )
+    # Wait for process to complete and gather stdout/stderr
+    return_code = await cmd_process.wait()
+    stdout = await stdout_task
+    stderr = await stderr_task
+    return CmdResult(stdout, stderr), return_code
+
+
+def kill_pid(pid: int, print_method: Callable[..., None] | None = None):
+    actual_print_method = print_method if print_method is not None else print
+    parent = psutil.Process(pid)
+    children = parent.children(recursive=True)
+    for child in children:
+        actual_print_method(f"Killing child process {child.pid}")
+        child.terminate()
+    actual_print_method(f"Killing process {pid}")
+    parent.terminate()

zrb/util/codemod/modification_mode.py

@@ -0,0 +1,3 @@
+PREPEND = 0
+APPEND = 1
+REPLACE = 2

zrb/util/codemod/modify_class.py

@@ -0,0 +1,58 @@
+import libcst as cst
+
+from zrb.util.codemod.modification_mode import APPEND, PREPEND, REPLACE
+
+
+def replace_class_code(original_code: str, class_name: str, new_code: str) -> str:
+    return _modify_code(original_code, class_name, new_code, REPLACE)
+
+
+def prepend_code_to_class(original_code: str, class_name: str, new_code: str) -> str:
+    return _modify_code(original_code, class_name, new_code, PREPEND)
+
+
+def append_code_to_class(original_code: str, class_name: str, new_code: str) -> str:
+    return _modify_code(original_code, class_name, new_code, APPEND)
+
+
+def _modify_code(original_code: str, class_name: str, new_code: str, mode: int) -> str:
+    # Parse the original code into a module
+    module = cst.parse_module(original_code)
+    # Initialize transformer with the class name and method code
+    transformer = _ClassCodeModifier(class_name, new_code, mode)
+    # Apply the transformation
+    modified_module = module.visit(transformer)
+    # Check if the class was found
+    if not transformer.class_found:
+        raise ValueError(f"Class {class_name} not found in the provided code.")
+    # Return the modified code
+    return modified_module.code
+
+
+class _ClassCodeModifier(cst.CSTTransformer):
+    def __init__(self, class_name: str, new_code: str, mode: int):
+        self.class_name = class_name
+        self.new_code = cst.parse_module(new_code).body
+        self.class_found = False
+        self.mode = mode
+
+    def leave_ClassDef(
+        self, original_node: cst.ClassDef, updated_node: cst.ClassDef
+    ) -> cst.ClassDef:
+        # Check if this is the target class
+        if original_node.name.value == self.class_name:
+            self.class_found = True
+            if self.mode == REPLACE:
+                new_body = updated_node.body.with_changes(body=tuple(self.new_code))
+                return updated_node.with_changes(body=new_body)
+            if self.mode == PREPEND:
+                new_body = updated_node.body.with_changes(
+                    body=tuple(self.new_code) + updated_node.body.body
+                )
+                return updated_node.with_changes(body=new_body)
+            if self.mode == APPEND:
+                new_body = updated_node.body.with_changes(
+                    body=updated_node.body.body + tuple(self.new_code)
+                )
+                return updated_node.with_changes(body=new_body)
+        return updated_node

zrb/util/codemod/modify_class_parent.py

@@ -0,0 +1,68 @@
+import libcst as cst
+
+from zrb.util.codemod.modification_mode import APPEND, PREPEND, REPLACE
+
+
+def replace_parent_class(
+    original_code: str, class_name: str, parent_class_name: str
+) -> str:
+    return _modify_parent_class(original_code, class_name, parent_class_name, REPLACE)
+
+
+def append_parent_class(
+    original_code: str, class_name: str, parent_class_name: str
+) -> str:
+    return _modify_parent_class(original_code, class_name, parent_class_name, APPEND)
+
+
+def prepend_parent_class(
+    original_code: str, class_name: str, parent_class_name: str
+) -> str:
+    return _modify_parent_class(original_code, class_name, parent_class_name, PREPEND)
+
+
+def _modify_parent_class(
+    original_code: str, class_name: str, parent_class_name: str, mode: int
+) -> str:
+    # Parse the original code into a module
+    module = cst.parse_module(original_code)
+    # Initialize transformer with the class name and parent class name
+    transformer = _ParentClassAdder(class_name, parent_class_name, mode)
+    # Apply the transformation
+    modified_module = module.visit(transformer)
+    # Check if the class was found
+    if not transformer.class_found:
+        raise ValueError(f"Class {class_name} not found in the provided code.")
+    # Return the modified code
+    return modified_module.code
+
+
+class _ParentClassAdder(cst.CSTTransformer):
+    def __init__(self, class_name: str, parent_class_name: str, mode: int):
+        self.class_name = class_name
+        self.parent_class_name = parent_class_name
+        self.class_found = False
+        self.mode = mode
+
+    def leave_ClassDef(
+        self, original_node: cst.ClassDef, updated_node: cst.ClassDef
+    ) -> cst.ClassDef:
+        # Check if this is the target class
+        if original_node.name.value == self.class_name:
+            self.class_found = True
+            if self.mode == REPLACE:
+                new_bases = (cst.Arg(value=cst.Name(self.parent_class_name)),)
+                return updated_node.with_changes(bases=new_bases)
+            if self.mode == PREPEND:
+                new_bases = (
+                    cst.Arg(value=cst.Name(self.parent_class_name)),
+                    *updated_node.bases,
+                )
+                return updated_node.with_changes(bases=new_bases)
+            if self.mode == APPEND:
+                new_bases = (
+                    *updated_node.bases,
+                    cst.Arg(value=cst.Name(self.parent_class_name)),
+                )
+                return updated_node.with_changes(bases=new_bases)
+        return updated_node

zrb/util/codemod/modify_class_property.py

@@ -0,0 +1,128 @@
+import libcst as cst
+
+from zrb.util.codemod.modification_mode import APPEND, PREPEND
+
+
+def append_property_to_class(
+    original_code: str,
+    class_name: str,
+    property_name: str,
+    annotation: str,
+    default_value: str,
+) -> str:
+    return _modify_class_property(
+        original_code, class_name, property_name, annotation, default_value, APPEND
+    )
+
+
+def prepend_property_to_class(
+    original_code: str,
+    class_name: str,
+    property_name: str,
+    annotation: str,
+    default_value: str,
+) -> str:
+    return _modify_class_property(
+        original_code, class_name, property_name, annotation, default_value, PREPEND
+    )
+
+
+def _modify_class_property(
+    original_code: str,
+    class_name: str,
+    property_name: str,
+    annotation: str,
+    default_value: str,
+    mode: int,
+) -> str:
+    # Parse the original code into a module
+    module = cst.parse_module(original_code)
+    # Initialize transformer with the class name, property name, annotation, and default value
+    transformer = _ClassPropertyModifier(
+        class_name, property_name, annotation, default_value, mode
+    )
+    # Apply the transformation
+    modified_module = module.visit(transformer)
+    # Check if the class was found
+    if not transformer.class_found:
+        raise ValueError(f"Class {class_name} not found in the provided code.")
+    # Return the modified code
+    return modified_module.code
+
+
+class _ClassPropertyModifier(cst.CSTTransformer):
+    def __init__(
+        self,
+        class_name: str,
+        property_name: str,
+        annotation: str,
+        default_value: str,
+        mode: int,
+    ):
+        self.class_name = class_name
+        self.property_name = property_name
+        self.annotation = cst.Annotation(cst.parse_expression(annotation))
+        self.default_value = cst.parse_expression(default_value)
+        self.class_found = False
+        self.mode = mode
+
+    def leave_ClassDef(
+        self, original_node: cst.ClassDef, updated_node: cst.ClassDef
+    ) -> cst.ClassDef:
+        # Check if this is the target class
+        if original_node.name.value == self.class_name:
+            self.class_found = True
+            # Create the annotated property with a default value
+            new_property = cst.SimpleStatementLine(
+                body=[
+                    cst.AnnAssign(
+                        target=cst.Name(self.property_name),
+                        annotation=self.annotation,
+                        value=self.default_value,
+                    )
+                ]
+            )
+            if self.mode == PREPEND:
+                new_body = cst.IndentedBlock(
+                    body=(new_property,) + updated_node.body.body
+                )
+                return updated_node.with_changes(body=new_body)
+            if self.mode == APPEND:
+                # Identify properties and methods
+                properties = []
+                methods = []
+                for stmt in updated_node.body.body:
+                    if isinstance(stmt, cst.SimpleStatementLine) and isinstance(
+                        stmt.body[0], (cst.AnnAssign, cst.Assign)
+                    ):
+                        properties.append(stmt)
+                    elif isinstance(stmt, cst.FunctionDef):
+                        methods.append(stmt)
+                if properties:
+                    # Class has properties
+                    last_property_index = updated_node.body.body.index(properties[-1])
+                    new_body = cst.IndentedBlock(
+                        body=(
+                            updated_node.body.body[: last_property_index + 1]
+                            + (new_property,)
+                            + updated_node.body.body[last_property_index + 1 :]
+                        )
+                    )
+                    return updated_node.with_changes(body=new_body)
+                if methods:
+                    # Class doesn't have properties but has methods
+                    first_method_index = updated_node.body.body.index(methods[0])
+                    new_body = cst.IndentedBlock(
+                        body=(
+                            updated_node.body.body[:first_method_index]
+                            + (new_property,)
+                            + updated_node.body.body[first_method_index:]
+                        )
+                    )
+                    return updated_node.with_changes(body=new_body)
+                # Class is empty, add add the bottom
+                new_body = cst.IndentedBlock(
+                    body=updated_node.body.body + (new_property,)
+                )
+                return updated_node.with_changes(body=new_body)
+        return updated_node

zrb/util/codemod/modify_dict.py

@@ -0,0 +1,75 @@
+import libcst as cst
+
+from zrb.util.codemod.modification_mode import APPEND, PREPEND
+
+
+def prepend_key_to_dict(
+    original_code: str, dictionary_name: str, new_key: str, new_value: str
+) -> str:
+    return _modify_dict(original_code, dictionary_name, new_key, new_value, PREPEND)
+
+
+def append_key_to_dict(
+    original_code: str, dictionary_name: str, new_key: str, new_value: str
+) -> str:
+    return _modify_dict(original_code, dictionary_name, new_key, new_value, APPEND)
+
+
+def _modify_dict(
+    original_code: str, dictionary_name: str, new_key: str, new_value: str, mode: int
+) -> str:
+    # Parse the original code into a module
+    module = cst.parse_module(original_code)
+    # Initialize the transformer with the necessary information
+    transformer = _DictionaryModifier(dictionary_name, new_key, new_value, mode)
+    # Apply the transformation
+    modified_module = module.visit(transformer)
+    # Error handling: raise an error if the dictionary is not found
+    if not transformer.found:
+        raise ValueError(
+            f"Dictionary {dictionary_name} not found in the provided code."
+        )
+    # Return the modified code
+    return modified_module.code
+
+
+class _DictionaryModifier(cst.CSTTransformer):
+    def __init__(self, dictionary_name: str, new_key: str, new_value: str, mode: int):
+        self.dictionary_name = dictionary_name
+        self.new_key = new_key
+        self.new_value = new_value
+        self.found = False
+        self.mode = mode
+
+    def leave_Assign(
+        self, original_node: cst.Assign, updated_node: cst.Assign
+    ) -> cst.Assign:
+        # Extract the first target from updated_node, which will be an AssignTarget
+        target = updated_node.targets[0]
+        # Check if the target is a Name (which should represent the dictionary)
+        if (
+            isinstance(target.target, cst.Name)
+            and target.target.value == self.dictionary_name
+        ):
+            # Check if it's a dictionary initialization (e.g., my_dict = {...})
+            if isinstance(updated_node.value, cst.Dict):
+                self.found = True
+                if self.mode == PREPEND:
+                    new_entries = (
+                        cst.DictElement(
+                            key=cst.SimpleString(f'"{self.new_key}"'),
+                            value=cst.SimpleString(f'"{self.new_value}"'),
+                        ),
+                    ) + updated_node.value.elements
+                    new_dict = updated_node.value.with_changes(elements=new_entries)
+                    return updated_node.with_changes(value=new_dict)
+                if self.mode == APPEND:
+                    new_entries = updated_node.value.elements + (
+                        cst.DictElement(
+                            key=cst.SimpleString(f'"{self.new_key}"'),
+                            value=cst.SimpleString(f'"{self.new_value}"'),
+                        ),
+                    )
+                    new_dict = updated_node.value.with_changes(elements=new_entries)
+                    return updated_node.with_changes(value=new_dict)
+        return updated_node

zrb/util/codemod/modify_function.py

@@ -0,0 +1,65 @@
+import libcst as cst
+
+from zrb.util.codemod.modification_mode import APPEND, PREPEND, REPLACE
+
+
+def replace_function_code(original_code: str, function_name: str, new_code: str) -> str:
+    return _modify_function(original_code, function_name, new_code, REPLACE)
+
+
+def prepend_code_to_function(
+    original_code: str, function_name: str, new_code: str
+) -> str:
+    return _modify_function(original_code, function_name, new_code, PREPEND)
+
+
+def append_code_to_function(
+    original_code: str, function_name: str, new_code: str
+) -> str:
+    return _modify_function(original_code, function_name, new_code, APPEND)
+
+
+def _modify_function(
+    original_code: str, function_name: str, new_code: str, mode: int
+) -> str:
+    # Parse the original code into a module
+    module = cst.parse_module(original_code)
+    # Initialize the transformer with the necessary information
+    transformer = _FunctionCodeModifier(function_name, new_code, mode)
+    # Apply the transformation
+    modified_module = module.visit(transformer)
+    # Error handling: raise an error if the class or function is not found
+    if not transformer.function_found:
+        raise ValueError(f"Function {function_name} not found.")
+    # Return the modified code
+    return modified_module.code
+
+
+class _FunctionCodeModifier(cst.CSTTransformer):
+    def __init__(self, function_name: str, new_code: str, mode: int):
+        self.function_name = function_name
+        # Use parse_module to handle multiple statements
+        self.new_code = cst.parse_module(new_code).body
+        self.function_found = False
+        self.mode = mode
+
+    def leave_FunctionDef(
+        self, original_node: cst.ClassDef, updated_node: cst.ClassDef
+    ) -> cst.ClassDef:
+        # Check if the class matches the target class
+        if original_node.name.value == self.function_name:
+            self.function_found = True
+            if self.mode == REPLACE:
+                new_body = updated_node.body.with_changes(body=tuple(self.new_code))
+                return updated_node.with_changes(body=new_body)
+            if self.mode == PREPEND:
+                new_body = updated_node.body.with_changes(
+                    body=tuple(self.new_code) + updated_node.body.body
+                )
+                return updated_node.with_changes(body=new_body)
+            if self.mode == APPEND:
+                new_body = updated_node.body.with_changes(
+                    body=updated_node.body.body + tuple(self.new_code)
+                )
+                return updated_node.with_changes(body=new_body)
+        return updated_node

zrb/util/codemod/modify_function_call.py

@@ -0,0 +1,68 @@
+import libcst as cst
+
+from zrb.util.codemod.modification_mode import APPEND, PREPEND, REPLACE
+
+
+def replace_function_call_param(
+    original_code: str, function_name: str, new_param: str
+) -> str:
+    return _modify_function_call(original_code, function_name, new_param, REPLACE)
+
+
+def prepend_param_to_function_call(
+    original_code: str, function_name: str, new_param: str
+) -> str:
+    return _modify_function_call(original_code, function_name, new_param, PREPEND)
+
+
+def append_param_to_function_call(
+    original_code: str, function_name: str, new_param: str
+) -> str:
+    return _modify_function_call(original_code, function_name, new_param, APPEND)
+
+
+def _modify_function_call(
+    original_code: str, function_name: str, new_param: str, mode: int
+) -> str:
+    # Parse the original code into a module
+    module = cst.parse_module(original_code)
+    # Initialize the transformer with the necessary information
+    transformer = _FunctionCallParamModifier(function_name, new_param, mode)
+    # Apply the transformation
+    modified_module = module.visit(transformer)
+    # Error handling: raise an error if the function call is not found
+    if not transformer.param_added:
+        raise ValueError(
+            f"Function call to {function_name} not found in the provided code."
+        )
+    # Return the modified code
+    return modified_module.code
+
+
+class _FunctionCallParamModifier(cst.CSTTransformer):
+    def __init__(self, func_name: str, new_param: str, mode: int):
+        self.func_name = func_name
+        # Parse the new parameter to ensure it’s a valid CST node
+        self.new_param = cst.parse_expression(new_param)
+        self.param_added = False
+        self.mode = mode
+
+    def leave_Call(self, original_node: cst.Call, updated_node: cst.Call) -> cst.Call:
+        # Check if the function call name matches the target function
+        if (
+            isinstance(original_node.func, cst.Name)
+            and original_node.func.value == self.func_name
+        ):
+            if self.mode == REPLACE:
+                new_args = (cst.Arg(value=self.new_param),)
+                self.param_added = True
+                return updated_node.with_changes(args=new_args)
+            if self.mode == PREPEND:
+                new_args = (cst.Arg(value=self.new_param),) + updated_node.args
+                self.param_added = True
+                return updated_node.with_changes(args=new_args)
+            if self.mode == APPEND:
+                new_args = updated_node.args + (cst.Arg(value=self.new_param),)
+                self.param_added = True
+                return updated_node.with_changes(args=new_args)
+        return updated_node