zrb 1.0.0a12__py3-none-any.whl → 1.0.0a15__py3-none-any.whl

zrb/session/session.py

@@ -28,6 +28,7 @@ from zrb.util.cli.style import (
     ICONS,
     MAGENTA,
     YELLOW,
+    remove_style,
 )
 from zrb.util.group import get_node_path
 from zrb.util.string.name import get_random_name
@@ -165,7 +166,9 @@ class Session(AnySession):
             main_task_name=self._main_task.name,
             path=self.task_path,
             final_result=(
-                f"{self.final_result}" if self.final_result is not None else ""
+                remove_style(f"{self.final_result}")
+                if self.final_result is not None
+                else ""
             ),
             finished=self.is_terminated,
             log=self.shared_ctx.shared_log,

zrb/task/cmd_task.py

@@ -1,6 +1,4 @@
-import asyncio
 import os
-import sys
 
 from zrb.attr.type import BoolAttr, IntAttr, StrAttr
 from zrb.cmd.cmd_result import CmdResult
@@ -12,7 +10,7 @@ from zrb.input.any_input import AnyInput
 from zrb.task.any_task import AnyTask
 from zrb.task.base_task import BaseTask
 from zrb.util.attr import get_int_attr, get_str_attr
-from zrb.util.cmd.command import check_unrecommended_commands
+from zrb.util.cmd.command import check_unrecommended_commands, run_command
 from zrb.util.cmd.remote import get_remote_cmd_script
 
 
@@ -119,43 +117,25 @@ class CmdTask(BaseTask):
         ctx.log_debug(f"Working directory: {cwd}")
         env_map = self.__get_env_map(ctx)
         ctx.log_debug(f"Environment map: {env_map}")
-        cmd_process = None
         if self._get_should_warn_unrecommended_commands():
             self._check_unrecommended_commands(ctx, shell, cmd_script)
-        try:
-            ctx.log_info("Running script")
-            cmd_process = await asyncio.create_subprocess_exec(
-                shell,
-                shell_flag,
-                cmd_script,
-                cwd=cwd,
-                stdin=sys.stdin if sys.stdin.isatty() else None,
-                stdout=asyncio.subprocess.PIPE,
-                stderr=asyncio.subprocess.PIPE,
-                env=env_map,
-                bufsize=0,
-            )
-            stdout_task = asyncio.create_task(
-                self.__read_stream(cmd_process.stdout, ctx.print, self._max_output_line)
-            )
-            stderr_task = asyncio.create_task(
-                self.__read_stream(cmd_process.stderr, ctx.print, self._max_error_line)
+        ctx.log_info("Running script")
+        cmd_result, return_code = await run_command(
+            cmd=[shell, shell_flag, cmd_script],
+            cwd=cwd,
+            env_map=env_map,
+            log_method=ctx.print,
+            max_output_line=self._max_output_line,
+            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.stderr}"
             )
-            # Wait for process to complete and gather stdout/stderr
-            return_code = await cmd_process.wait()
-            stdout = await stdout_task
-            stderr = await stderr_task
-            # Check for errors
-            if return_code != 0:
-                ctx.log_error(f"Exit status: {return_code}")
-                raise Exception(
-                    f"Process {self._name} exited ({return_code}): {stderr}"
-                )
-            ctx.log_info(f"Exit status: {return_code}")
-            return CmdResult(stdout, stderr)
-        finally:
-            if cmd_process is not None and cmd_process.returncode is None:
-                cmd_process.terminate()
+        ctx.log_info(f"Exit status: {return_code}")
+        return cmd_result
 
     def _get_should_warn_unrecommended_commands(self):
         if self._should_warn_unrecommended_command is None:
@@ -178,19 +158,6 @@ class CmdTask(BaseTask):
         envs["PYTHONBUFFERED"] = "1"
         return envs
 
-    async def __read_stream(self, stream, log_method, max_lines):
-        lines = []
-        while True:
-            line = await stream.readline()
-            if not line:
-                break
-            line = line.decode("utf-8").rstrip()
-            lines.append(line)
-            if len(lines) > max_lines:
-                lines.pop(0)  # Keep only the last max_lines
-            log_method(line)
-        return "\n".join(lines)
-
     def _get_shell(self, ctx: AnyContext) -> str:
         return get_str_attr(
             ctx, self._shell, DEFAULT_SHELL, auto_render=self._render_shell

zrb/task/scaffolder.py

@@ -92,8 +92,8 @@ class Scaffolder(BaseTask):
         destination_path = self._get_destination_path(ctx)
         self._copy_path(ctx, source_path, destination_path)
         transformers = self._get_content_transformers()
-        file_paths = self._get_all_file_paths(destination_path)
-        for file_path in file_paths:
+        file_path_list = self._get_all_file_paths(destination_path)
+        for file_path in file_path_list:
             for transformer in transformers:
                 if transformer.match(ctx, file_path):
                     try:
@@ -118,14 +118,10 @@ class Scaffolder(BaseTask):
                         dest_dir, self._transform_path(ctx, file_name)
                     )
                     shutil.copy2(src_file, dest_file)
-                    ctx.log_info(f"Copied and renamed {src_file} to {dest_file}")
+                    ctx.log_info(f"Copied {src_file} to {dest_file}")
         else:
-            dest_file = os.path.join(
-                destination_path,
-                self._transform_path(ctx, os.path.basename(source_path)),
-            )
-            shutil.copy2(source_path, dest_file)
-            ctx.log_info(f"Copied and renamed {source_path} to {dest_file}")
+            shutil.copy2(source_path, destination_path)
+            ctx.log_info(f"Copied {source_path} to {destination_path}")
 
     def _transform_path(self, ctx: AnyContext, file_path: str):
         if callable(self._path_transformer):
@@ -139,6 +135,8 @@ class Scaffolder(BaseTask):
         """
         Returns a list of absolute file paths for all files in the given path, recursively.
         """
+        if os.path.isfile(path):
+            return [os.path.abspath(path)]
         file_paths = []
         for root, _, files in os.walk(path):
             for file in files:

zrb/util/cli/style.py

@@ -1,3 +1,5 @@
+import re
+
 BLACK = 30
 RED = 31
 GREEN = 32
@@ -118,6 +120,11 @@ ICONS = [
 ]
 
 
+def remove_style(text):
+    ansi_escape = re.compile(r"\x1B[@-_][0-?]*[ -/]*[@-~]")
+    return ansi_escape.sub("", text)
+
+
 def stylize(
     text: str,
     color: int | None = None,

zrb/util/cmd/command.py

@@ -1,4 +1,10 @@
+import asyncio
+import os
 import re
+import sys
+from collections.abc import Callable
+
+from zrb.cmd.cmd_result import CmdResult
 
 
 def check_unrecommended_commands(cmd_script: str) -> dict[str, str]:
@@ -31,3 +37,57 @@ def check_unrecommended_commands(cmd_script: str) -> dict[str, str]:
         if re.search(pattern, cmd_script):
             violations[pattern] = reason
     return violations
+
+
+async def run_command(
+    cmd: list[str],
+    cwd: str | None = None,
+    env_map: dict[str, str] | None = None,
+    log_method: Callable[..., None] = print,
+    max_output_line: int = 1000,
+    max_error_line: int = 1000,
+) -> tuple[CmdResult, int]:
+
+    async def __read_stream(
+        stream, log_method: Callable[..., None], max_lines: int
+    ) -> str:
+        lines = []
+        while True:
+            line = await stream.readline()
+            if not line:
+                break
+            line = line.decode("utf-8").rstrip()
+            lines.append(line)
+            if len(lines) > max_lines:
+                lines.pop(0)  # Keep only the last max_lines
+            log_method(line)
+        return "\n".join(lines)
+
+    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,
+        )
+        stdout_task = asyncio.create_task(
+            __read_stream(cmd_process.stdout, log_method, max_output_line)
+        )
+        stderr_task = asyncio.create_task(
+            __read_stream(cmd_process.stderr, log_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 is not None and cmd_process.returncode is None:
+            cmd_process.terminate()

zrb/util/codemod/add_code_to_module.py

@@ -0,0 +1,12 @@
+def add_code_to_module(source_code: str, new_code: str) -> str:
+    lines = source_code.splitlines()
+    last_import_index = -1
+    for i, line in enumerate(lines):
+        stripped_line = line.strip()
+        if stripped_line.startswith("import") or stripped_line.startswith("from"):
+            last_import_index = i
+        elif stripped_line and not stripped_line.startswith("#"):
+            break
+    if last_import_index != -1:
+        lines.insert(last_import_index + 1, new_code)
+    return "\n".join(lines)

zrb/util/git.py

@@ -1,8 +1,11 @@
 import os
-import subprocess
+from collections.abc import Callable
+from typing import Any
 
 from pydantic import BaseModel
 
+from zrb.util.cmd.command import run_command
+
 
 class DiffResult(BaseModel):
     created: list[str]
@@ -10,19 +13,20 @@ class DiffResult(BaseModel):
     updated: list[str]
 
 
-def get_diff(repo_dir: str, source_commit: str, current_commit: str) -> DiffResult:
-    try:
-        result = subprocess.run(
-            ["git", "diff", source_commit, current_commit],
-            stdout=subprocess.PIPE,
-            stderr=subprocess.PIPE,
-            cwd=repo_dir,
-            text=True,
-            check=True,
-        )
-    except subprocess.CalledProcessError as e:
-        raise Exception(e.stderr or e.stdout)
-    lines = result.stdout.strip().split("\n")
+async def get_diff(
+    repo_dir: str,
+    source_commit: str,
+    current_commit: str,
+    log_method: Callable[..., Any] = print,
+) -> DiffResult:
+    cmd_result, exit_code = await run_command(
+        cmd=["git", "diff", source_commit, current_commit],
+        cwd=repo_dir,
+        log_method=log_method,
+    )
+    if exit_code != 0:
+        raise Exception(f"Non zero exit code: {exit_code}")
+    lines = cmd_result.output.strip().split("\n")
     diff: dict[str, dict[str, bool]] = {}
     for line in lines:
         if not line.startswith("---") and not line.startswith("+++"):
@@ -50,125 +54,103 @@ def get_diff(repo_dir: str, source_commit: str, current_commit: str) -> DiffResu
     )
 
 
-def get_repo_dir() -> str:
-    try:
-        # Run the Git command to get the repository's top-level directory
-        result = subprocess.run(
-            ["git", "rev-parse", "--show-toplevel"],
-            stdout=subprocess.PIPE,
-            stderr=subprocess.PIPE,
-            text=True,
-            check=True,
-        )
-        # Return the directory path
-        return os.path.abspath(result.stdout.strip())
-    except subprocess.CalledProcessError as e:
-        raise Exception(e.stderr or e.stdout)
-
-
-def get_current_branch(repo_dir: str) -> str:
-    try:
-        result = subprocess.run(
-            ["git", "rev-parse", "--abbrev-ref", "HEAD"],
-            stdout=subprocess.PIPE,
-            stderr=subprocess.PIPE,
-            cwd=repo_dir,
-            text=True,
-            check=True,
-        )
-        return result.stdout.strip()
-    except subprocess.CalledProcessError as e:
-        raise Exception(e.stderr or e.stdout)
-
-
-def get_branches(repo_dir: str) -> list[str]:
-    try:
-        result = subprocess.run(
-            ["git", "branch"],
-            stdout=subprocess.PIPE,
-            stderr=subprocess.PIPE,
-            cwd=repo_dir,
-            text=True,
-            check=True,
-        )
-        return [
-            branch.lstrip("*").strip() for branch in result.stdout.strip().split("\n")
-        ]
-    except subprocess.CalledProcessError as e:
-        raise Exception(e.stderr or e.stdout)
-
-
-def delete_branch(repo_dir: str, branch_name: str) -> str:
-    try:
-        result = subprocess.run(
-            ["git", "branch", "-D", branch_name],
-            stdout=subprocess.PIPE,
-            stderr=subprocess.PIPE,
-            cwd=repo_dir,
-            text=True,
-            check=True,
-        )
-        return result.stdout.strip()
-    except subprocess.CalledProcessError as e:
-        raise Exception(e.stderr or e.stdout)
-
-
-def add(repo_dir: str) -> str:
-    try:
-        subprocess.run(
-            ["git", "add", ".", "-A"],
-            stdout=subprocess.PIPE,
-            stderr=subprocess.PIPE,
-            cwd=repo_dir,
-            text=True,
-            check=True,
-        )
-    except subprocess.CalledProcessError as e:
-        raise Exception(e.stderr or e.stdout)
-
-
-def commit(repo_dir: str, message: str) -> str:
-    try:
-        subprocess.run(
-            ["git", "commit", "-m", message],
-            stdout=subprocess.PIPE,
-            stderr=subprocess.PIPE,
-            cwd=repo_dir,
-            text=True,
-            check=True,
-        )
-    except subprocess.CalledProcessError as e:
+async def get_repo_dir(log_method: Callable[..., Any] = print) -> str:
+    cmd_result, exit_code = await run_command(
+        cmd=["git", "rev-parse", "--show-toplevel"],
+        log_method=log_method,
+    )
+    if exit_code != 0:
+        raise Exception(f"Non zero exit code: {exit_code}")
+    return os.path.abspath(cmd_result.output.strip())
+
+
+async def get_current_branch(
+    repo_dir: str, log_method: Callable[..., Any] = print
+) -> str:
+    cmd_result, exit_code = await run_command(
+        cmd=["git", "rev-parse", "--abbrev-ref", "HEAD"],
+        cwd=repo_dir,
+        log_method=log_method,
+    )
+    if exit_code != 0:
+        raise Exception(f"Non zero exit code: {exit_code}")
+    return cmd_result.output.strip()
+
+
+async def get_branches(
+    repo_dir: str, log_method: Callable[..., Any] = print
+) -> list[str]:
+    cmd_result, exit_code = await run_command(
+        cmd=["git", "rev-parse", "--abbrev-ref", "HEAD"],
+        cwd=repo_dir,
+        log_method=log_method,
+    )
+    if exit_code != 0:
+        raise Exception(f"Non zero exit code: {exit_code}")
+    return [
+        branch.lstrip("*").strip() for branch in cmd_result.output.strip().split("\n")
+    ]
+
+
+async def delete_branch(
+    repo_dir: str, branch_name: str, log_method: Callable[..., Any] = print
+) -> str:
+    cmd_result, exit_code = await run_command(
+        cmd=["git", "branch", "-D", branch_name],
+        cwd=repo_dir,
+        log_method=log_method,
+    )
+    if exit_code != 0:
+        raise Exception(f"Non zero exit code: {exit_code}")
+    return cmd_result.output.strip()
+
+
+async def add(repo_dir: str, log_method: Callable[..., Any] = print):
+    _, exit_code = await run_command(
+        cmd=["git", "add", ".", "-A"],
+        cwd=repo_dir,
+        log_method=log_method,
+    )
+    if exit_code != 0:
+        raise Exception(f"Non zero exit code: {exit_code}")
+
+
+async def commit(
+    repo_dir: str, message: str, log_method: Callable[..., Any] = print
+) -> str:
+    cmd_result, exit_code = await run_command(
+        cmd=["git", "commit", "-m", message],
+        cwd=repo_dir,
+        log_method=log_method,
+    )
+    if exit_code != 0:
         ignored_error_message = "nothing to commit, working tree clean"
         if (
-            ignored_error_message not in e.stderr
-            and ignored_error_message not in e.stdout
+            ignored_error_message not in cmd_result.error
+            and ignored_error_message not in cmd_result.output
         ):
-            raise Exception(e.stderr or e.stdout)
-
-
-def pull(repo_dir: str, remote: str, branch: str) -> str:
-    try:
-        subprocess.run(
-            ["git", "pull", remote, branch],
-            stdout=subprocess.PIPE,
-            stderr=subprocess.PIPE,
-            cwd=repo_dir,
-            text=True,
-            check=True,
-        )
-    except subprocess.CalledProcessError as e:
-        raise Exception(e.stderr or e.stdout)
-
-
-def push(repo_dir: str, remote: str, branch: str) -> str:
-    try:
-        subprocess.run(
-            ["git", "push", "-u", remote, branch],
-            stdout=subprocess.PIPE,
-            stderr=subprocess.PIPE,
-            cwd=repo_dir,
-            text=True,
-            check=True,
-        )
-    except subprocess.CalledProcessError as e:
-        raise Exception(e.stderr or e.stdout)
+            raise Exception(f"Non zero exit code: {exit_code}")
+
+
+async def pull(
+    repo_dir: str, remote: str, branch: str, log_method: Callable[..., Any] = print
+) -> str:
+    _, exit_code = await run_command(
+        cmd=["git", "pull", remote, branch],
+        cwd=repo_dir,
+        log_method=log_method,
+    )
+    if exit_code != 0:
+        raise Exception(f"Non zero exit code: {exit_code}")
+
+
+async def push(
+    repo_dir: str, remote: str, branch: str, log_method: Callable[..., Any] = print
+) -> str:
+    _, exit_code = await run_command(
+        cmd=["git", "push", "-u", remote, branch],
+        cwd=repo_dir,
+        log_method=log_method,
+    )
+    if exit_code != 0:
+        raise Exception(f"Non zero exit code: {exit_code}")