zrb 0.28.1__py3-none-any.whl → 1.0.0__py3-none-any.whl

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

zrb/util/codemod/modify_method.py

@@ -0,0 +1,88 @@
+import libcst as cst
+
+from zrb.util.codemod.modification_mode import APPEND, PREPEND, REPLACE
+
+
+def replace_method_code(
+    original_code: str, class_name: str, method_name: str, new_code: str
+) -> str:
+    return _modify_method(original_code, class_name, method_name, new_code, REPLACE)
+
+
+def prepend_code_to_method(
+    original_code: str, class_name: str, method_name: str, new_code: str
+) -> str:
+    return _modify_method(original_code, class_name, method_name, new_code, PREPEND)
+
+
+def append_code_to_method(
+    original_code: str, class_name: str, method_name: str, new_code: str
+) -> str:
+    return _modify_method(original_code, class_name, method_name, new_code, APPEND)
+
+
+def _modify_method(
+    original_code: str, class_name: str, method_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 = _MethodModifier(class_name, method_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.class_found:
+        raise ValueError(f"Class {class_name} not found in the provided code.")
+    if not transformer.method_found:
+        raise ValueError(f"Method {method_name} not found in class {class_name}.")
+    # Return the modified code
+    return modified_module.code
+
+
+class _MethodModifier(cst.CSTTransformer):
+    def __init__(self, class_name: str, method_name: str, new_code: str, mode: int):
+        self.class_name = class_name
+        self.method_name = method_name
+        # Use parse_module to handle multiple statements
+        self.new_code = cst.parse_module(new_code).body
+        self.class_found = False
+        self.method_found = False
+        self.mode = mode
+
+    def leave_ClassDef(
+        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.class_name:
+            self.class_found = True
+            # Now, modify function definitions inside this class
+            new_body = []
+            for (
+                item
+            ) in updated_node.body.body:  # Access body.body, not just updated_node.body
+                if (
+                    isinstance(item, cst.FunctionDef)
+                    and item.name.value == self.method_name
+                ):
+                    # Modify the target function by adding the new code
+                    if self.mode == REPLACE:
+                        body_with_new_code = item.body.with_changes(
+                            body=tuple(self.new_code)
+                        )
+                        new_body.append(item.with_changes(body=body_with_new_code))
+                    if self.mode == PREPEND:
+                        body_with_new_code = item.body.with_changes(
+                            body=tuple(self.new_code) + item.body.body
+                        )
+                        new_body.append(item.with_changes(body=body_with_new_code))
+
+                    if self.mode == APPEND:
+                        body_with_new_code = item.body.with_changes(
+                            body=item.body.body + tuple(self.new_code)
+                        )
+                        new_body.append(item.with_changes(body=body_with_new_code))
+                    self.method_found = True
+                else:
+                    new_body.append(item)
+            return updated_node.with_changes(body=cst.IndentedBlock(new_body))
+        return updated_node

zrb/util/codemod/modify_module.py

@@ -0,0 +1,11 @@
+def prepend_code_to_module(original_code: str, new_code: str) -> str:
+    lines = original_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
+        else:
+            break
+    lines.insert(last_import_index + 1, new_code)
+    return "\n".join(lines)

zrb/util/cron.py

@@ -0,0 +1,71 @@
+import datetime
+
+
+def parse_cron_field(field: str, min_value: int, max_value: int):
+    """Parse a cron field with support for wildcards (*), ranges, lists, and steps."""
+    values = set()
+    if field == "*":
+        return set(range(min_value, max_value + 1))
+    elif field.startswith("*/"):
+        # Handle step values like "*/5"
+        step = int(field[2:])
+        return set(range(min_value, max_value + 1, step))
+    elif "-" in field or "," in field or "/" in field:
+        # Handle ranges, lists, and steps
+        parts = field.split(",")
+        for part in parts:
+            if "/" in part:
+                range_part, step = part.split("/")
+                step = int(step)
+                if "-" in range_part:
+                    start, end = map(int, range_part.split("-"))
+                else:
+                    start, end = int(range_part), max_value
+                values.update(range(start, end + 1, step))
+            elif "-" in part:
+                start, end = map(int, part.split("-"))
+                values.update(range(start, end + 1))
+            else:
+                values.add(int(part))
+    else:
+        # Handle individual values
+        values.add(int(field))
+    return values
+
+
+def handle_special_cron_patterns(pattern: str, dt: datetime.datetime):
+    """Handle special cron patterns like @yearly, @monthly, etc."""
+    if pattern == "@yearly" or pattern == "@annually":
+        return dt.month == 1 and dt.day == 1 and dt.hour == 0 and dt.minute == 0
+    elif pattern == "@monthly":
+        return dt.day == 1 and dt.hour == 0 and dt.minute == 0
+    elif pattern == "@weekly":
+        return (
+            dt.weekday() == 0 and dt.hour == 0 and dt.minute == 0
+        )  # Monday at midnight
+    elif pattern == "@daily" or pattern == "@midnight":
+        return dt.hour == 0 and dt.minute == 0
+    elif pattern == "@hourly":
+        return dt.minute == 0
+    elif pattern == "@minutely":
+        return True
+    return False
+
+
+def match_cron(cron_pattern: str, dt: datetime.datetime):
+    """Check if a datetime object matches a cron pattern, including special cases."""
+    # Handle special cron patterns
+    if cron_pattern.startswith("@"):
+        return handle_special_cron_patterns(cron_pattern, dt)
+    minute, hour, day, month, day_of_week = cron_pattern.split()
+    # Parse each field with step, range, and wildcard support
+    minutes = parse_cron_field(minute, 0, 59)
+    hours = parse_cron_field(hour, 0, 23)
+    days = parse_cron_field(day, 1, 31)
+    months = parse_cron_field(month, 1, 12)
+    days_of_week = parse_cron_field(day_of_week, 0, 6)
+    # Match with standard cron behavior (day OR day_of_week must match)
+    day_match = dt.day in days or dt.weekday() in days_of_week
+    return (
+        dt.minute in minutes and dt.hour in hours and dt.month in months and day_match
+    )

zrb/util/file.py

@@ -0,0 +1,25 @@
+import os
+import re
+
+
+def read_file(file_path: str, replace_map: dict[str, str] = {}) -> str:
+    with open(file_path, "r", encoding="utf-8") as f:
+        content = f.read()
+    for key, val in replace_map.items():
+        content = content.replace(key, val)
+    return content
+
+
+def write_file(file_path: str, content: str | list[str]):
+    if isinstance(content, list):
+        content = "\n".join([line for line in content if line is not None])
+    dir_path = os.path.dirname(file_path)
+    os.makedirs(dir_path, exist_ok=True)
+    should_add_eol = content.endswith("\n")
+    # Remove trailing newlines, but keep one if the file originally ended up with newline
+    content = re.sub(r"\n{3,}$", "\n\n", content)
+    content = content.rstrip("\n")
+    if should_add_eol:
+        content += "\n"
+    with open(file_path, "w") as f:
+        f.write(content)

zrb/util/git.py

@@ -0,0 +1,156 @@
+import os
+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]
+    removed: list[str]
+    updated: list[str]
+
+
+async def get_diff(
+    repo_dir: str,
+    source_commit: str,
+    current_commit: str,
+    print_method: Callable[..., Any] = print,
+) -> DiffResult:
+    cmd_result, exit_code = await run_command(
+        cmd=["git", "diff", source_commit, current_commit],
+        cwd=repo_dir,
+        print_method=print_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("+++"):
+            continue
+        if line[4:6] != "a/" and line[4:6] != "b/":
+            continue
+        # line should contains something like `--- a/some-file.txt`
+        file = line[6:]
+        if file not in diff:
+            diff[file] = {"plus": False, "minus": False}
+        if line.startswith("---"):
+            diff[file]["minus"] = True
+        if line.startswith("+++"):
+            diff[file]["plus"] = True
+    return DiffResult(
+        created=[
+            file for file, state in diff.items() if state["plus"] and not state["minus"]
+        ],
+        removed=[
+            file for file, state in diff.items() if not state["plus"] and state["minus"]
+        ],
+        updated=[
+            file for file, state in diff.items() if state["plus"] and state["minus"]
+        ],
+    )
+
+
+async def get_repo_dir(print_method: Callable[..., Any] = print) -> str:
+    cmd_result, exit_code = await run_command(
+        cmd=["git", "rev-parse", "--show-toplevel"],
+        print_method=print_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, print_method: Callable[..., Any] = print
+) -> str:
+    cmd_result, exit_code = await run_command(
+        cmd=["git", "rev-parse", "--abbrev-ref", "HEAD"],
+        cwd=repo_dir,
+        print_method=print_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, print_method: Callable[..., Any] = print
+) -> list[str]:
+    cmd_result, exit_code = await run_command(
+        cmd=["git", "rev-parse", "--abbrev-ref", "HEAD"],
+        cwd=repo_dir,
+        print_method=print_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, print_method: Callable[..., Any] = print
+) -> str:
+    cmd_result, exit_code = await run_command(
+        cmd=["git", "branch", "-D", branch_name],
+        cwd=repo_dir,
+        print_method=print_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, print_method: Callable[..., Any] = print):
+    _, exit_code = await run_command(
+        cmd=["git", "add", ".", "-A"],
+        cwd=repo_dir,
+        print_method=print_method,
+    )
+    if exit_code != 0:
+        raise Exception(f"Non zero exit code: {exit_code}")
+
+
+async def commit(
+    repo_dir: str, message: str, print_method: Callable[..., Any] = print
+) -> str:
+    cmd_result, exit_code = await run_command(
+        cmd=["git", "commit", "-m", message],
+        cwd=repo_dir,
+        print_method=print_method,
+    )
+    if exit_code != 0:
+        ignored_error_message = "nothing to commit, working tree clean"
+        if (
+            ignored_error_message not in cmd_result.error
+            and ignored_error_message not in cmd_result.output
+        ):
+            raise Exception(f"Non zero exit code: {exit_code}")
+
+
+async def pull(
+    repo_dir: str, remote: str, branch: str, print_method: Callable[..., Any] = print
+) -> str:
+    _, exit_code = await run_command(
+        cmd=["git", "pull", remote, branch],
+        cwd=repo_dir,
+        print_method=print_method,
+    )
+    if exit_code != 0:
+        raise Exception(f"Non zero exit code: {exit_code}")
+
+
+async def push(
+    repo_dir: str, remote: str, branch: str, print_method: Callable[..., Any] = print
+) -> str:
+    _, exit_code = await run_command(
+        cmd=["git", "push", "-u", remote, branch],
+        cwd=repo_dir,
+        print_method=print_method,
+    )
+    if exit_code != 0:
+        raise Exception(f"Non zero exit code: {exit_code}")

zrb/util/git_subtree.py

@@ -0,0 +1,112 @@
+import os
+from collections.abc import Callable
+from typing import Any
+
+from pydantic import BaseModel
+
+from zrb.util.cmd.command import run_command
+from zrb.util.file import read_file, write_file
+
+
+class SingleSubTreeConfig(BaseModel):
+    repo_url: str
+    branch: str
+    prefix: str
+
+
+class SubTreeConfig(BaseModel):
+    data: dict[str, SingleSubTreeConfig]
+
+
+def load_config(repo_dir: str) -> SubTreeConfig:
+    file_path = os.path.join(repo_dir, "subtrees.json")
+    if not os.path.exists(file_path):
+        return SubTreeConfig(data={})
+    return SubTreeConfig.model_validate_json(read_file(file_path))
+
+
+def save_config(repo_dir: str, config: SubTreeConfig):
+    file_path = os.path.join(repo_dir, "subtrees.json")
+    write_file(file_path, config.model_dump_json(indent=2))
+
+
+async def add_subtree(
+    repo_dir: str,
+    name: str,
+    repo_url: str,
+    branch: str,
+    prefix: str,
+    print_method: Callable[..., Any] = print,
+):
+    config = load_config()
+    if os.path.isdir(prefix):
+        raise ValueError(f"Directory exists: {prefix}")
+    if name in config.data:
+        raise ValueError(f"Subtree config already exists: {name}")
+    _, exit_code = await run_command(
+        cmd=[
+            "git",
+            "subtree",
+            "add",
+            "--prefix",
+            prefix,
+            repo_url,
+            branch,
+        ],
+        cwd=repo_dir,
+        print_method=print_method,
+    )
+    if exit_code != 0:
+        raise Exception(f"Non zero exit code: {exit_code}")
+    config.data[name] = SingleSubTreeConfig(
+        repo_url=repo_url, branch=branch, prefix=prefix
+    )
+    save_config(repo_dir, config)
+
+
+async def pull_subtree(
+    repo_dir: str,
+    prefix: str,
+    repo_url: str,
+    branch: str,
+    print_method: Callable[..., Any] = print,
+):
+    _, exit_code = await run_command(
+        cmd=[
+            "git",
+            "subtree",
+            "pull",
+            "--prefix",
+            prefix,
+            repo_url,
+            branch,
+        ],
+        cwd=repo_dir,
+        print_method=print_method,
+    )
+    if exit_code != 0:
+        raise Exception(f"Non zero exit code: {exit_code}")
+
+
+async def push_subtree(
+    repo_dir: str,
+    prefix: str,
+    repo_url: str,
+    branch: str,
+    print_method: Callable[..., Any] = print,
+):
+    _, exit_code = await run_command(
+        cmd=[
+            "git",
+            "subtree",
+            "push",
+            "--prefix",
+            prefix,
+            repo_url,
+            branch,
+        ],
+        cwd=repo_dir,
+        print_method=print_method,
+    )
+    if exit_code != 0:
+        raise Exception(f"Non zero exit code: {exit_code}")

zrb/util/group.py

@@ -0,0 +1,86 @@
+from zrb.group.any_group import AnyGroup
+from zrb.task.any_task import AnyTask
+
+
+class NodeNotFoundError(ValueError):
+    pass
+
+
+def extract_node_from_args(
+    root_group: AnyGroup, args: list[str], web_only: bool = False
+) -> tuple[AnyGroup | AnyTask, list[str], list[str]]:
+    node = root_group
+    node_path = []
+    residual_args = []
+    for index, name in enumerate(args):
+        task = node.get_task_by_alias(name)
+        if web_only and task is not None and task.cli_only:
+            task = None
+        group = node.get_group_by_alias(name)
+        if group is not None and len(get_all_subtasks(group, web_only)) == 0:
+            # If group doesn't contain any task, then ignore its existence
+            group = None
+        if task is None and group is None:
+            raise NodeNotFoundError(
+                f"Invalid subcommand: {root_group.name} {' '.join(args)}"
+            )
+        node_path.append(name)
+        if group is not None:
+            if task is not None and index == len(args) - 1:
+                node = task
+                residual_args = args[index + 1 :]
+                break
+            node = group
+            continue
+        if task is not None:
+            node = task
+            residual_args = args[index + 1 :]
+            break
+    return node, node_path, residual_args
+
+
+def get_node_path(group: AnyGroup, node: AnyGroup | AnyTask) -> list[str] | None:
+    if group is None:
+        return []
+    if isinstance(node, AnyTask):
+        for alias, subtask in group.subtasks.items():
+            if subtask == node:
+                return [alias]
+    if isinstance(node, AnyGroup):
+        for alias, subgroup in group.subgroups.items():
+            if subgroup == node:
+                return [alias]
+    for alias, subgroup in group.subgroups.items():
+        result = get_node_path(subgroup, node)
+        if result is not None:
+            return [alias] + result
+    return None
+
+
+def get_non_empty_subgroups(
+    group: AnyGroup, web_only: bool = False
+) -> dict[str, AnyGroup]:
+    return {
+        alias: subgroup
+        for alias, subgroup in group.subgroups.items()
+        if len(get_all_subtasks(subgroup, web_only)) > 0
+    }
+
+
+def get_subtasks(group: AnyGroup, web_only: bool = False) -> dict[str, AnyTask]:
+    return {
+        alias: subtask
+        for alias, subtask in group.subtasks.items()
+        if not web_only or (web_only and not subtask.cli_only)
+    }
+
+
+def get_all_subtasks(group: AnyGroup, web_only: bool = False) -> list[AnyTask]:
+    subtasks = [
+        subtask
+        for subtask in group.subtasks.values()
+        if not web_only or (web_only and not subtask.cli_only)
+    ]
+    for subgroup in group.subgroups.values():
+        subtasks += get_all_subtasks(subgroup, web_only)
+    return subtasks