zrb 1.0.0a2__py3-none-any.whl → 1.0.0a3__py3-none-any.whl

zrb/task/scheduler.py

@@ -9,11 +9,10 @@ from zrb.context.any_shared_context import AnySharedContext
 from zrb.env.any_env import AnyEnv
 from zrb.input.any_input import AnyInput
 from zrb.task.any_task import AnyTask
-
-from ..util.attr import get_str_attr
-from ..util.cron import match_cron
-from ..xcom.xcom import Xcom
-from .base_trigger import BaseTrigger
+from zrb.task.base_trigger import BaseTrigger
+from zrb.util.attr import get_str_attr
+from zrb.util.cron import match_cron
+from zrb.xcom.xcom import Xcom
 
 
 class Scheduler(BaseTrigger):
@@ -25,8 +24,8 @@ class Scheduler(BaseTrigger):
         icon: str | None = None,
         description: str | None = None,
         cli_only: bool = False,
-        input: list[AnyInput] | AnyInput | None = None,
-        env: list[AnyEnv] | AnyEnv | None = None,
+        input: list[AnyInput | None] | AnyInput | None = None,
+        env: list[AnyEnv | None] | AnyEnv | None = None,
         schedule: StrAttr = None,
         execute_condition: bool | str | Callable[[AnySharedContext], bool] = True,
         queue_name: fstring | None = None,

zrb/task/task.py

@@ -1,4 +1,4 @@
-from .base_task import BaseTask
+from zrb.task.base_task import BaseTask
 
 
 class Task(BaseTask):

zrb/task/tcp_check.py

@@ -1,14 +1,14 @@
 import asyncio
 from collections.abc import Callable
 
-from ..attr.type import IntAttr, StrAttr
-from ..context.any_context import AnyContext
-from ..context.context import Context
-from ..env.any_env import AnyEnv
-from ..input.any_input import AnyInput
-from ..util.attr import get_int_attr, get_str_attr
-from .any_task import AnyTask
-from .base_task import BaseTask
+from zrb.attr.type import IntAttr, StrAttr
+from zrb.context.any_context import AnyContext
+from zrb.context.context import Context
+from zrb.env.any_env import AnyEnv
+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
 
 
 class TcpCheck(BaseTask):

zrb/util/attr.py

@@ -1,8 +1,24 @@
 from typing import Any
 
-from ..attr.type import AnyAttr, BoolAttr, FloatAttr, IntAttr, StrAttr, StrDictAttr
-from ..context.any_shared_context import AnySharedContext
-from ..util.string.conversion import to_boolean
+from zrb.attr.type import (
+    AnyAttr,
+    BoolAttr,
+    FloatAttr,
+    IntAttr,
+    StrAttr,
+    StrDictAttr,
+    StrListAttr,
+)
+from zrb.context.any_shared_context import AnySharedContext
+from zrb.util.string.conversion import to_boolean
+
+
+def get_str_list_attr(
+    shared_ctx: AnySharedContext, attr: StrListAttr | None, auto_render: bool = True
+) -> list[str]:
+    if callable(attr):
+        return attr(shared_ctx)
+    return {get_str_attr(shared_ctx, val, "", auto_render) for val in attr}
 
 
 def get_str_dict_attr(

zrb/util/cli/style.py

@@ -6,8 +6,33 @@ BLUE = 34
 MAGENTA = 35
 CYAN = 36
 WHITE = 37
-
-VALID_COLORS = [BLACK, RED, GREEN, YELLOW, BLUE, MAGENTA, CYAN, WHITE]
+BRIGHT_BLACK = 90
+BRIGHT_RED = 91
+BRIGHT_GREEN = 92
+BRIGHT_YELLOW = 93
+BRIGHT_BLUE = 94
+BRIGHT_MAGENTA = 95
+BRIGHT_CYAN = 96
+BRIGHT_WHITE = 97
+
+VALID_COLORS = [
+    BLACK,
+    RED,
+    GREEN,
+    YELLOW,
+    BLUE,
+    MAGENTA,
+    CYAN,
+    WHITE,
+    BRIGHT_BLACK,
+    BRIGHT_RED,
+    BRIGHT_GREEN,
+    BRIGHT_YELLOW,
+    BRIGHT_BLUE,
+    BRIGHT_MAGENTA,
+    BRIGHT_CYAN,
+    BRIGHT_WHITE,
+]
 
 BG_BLACK = 40
 BG_RED = 41
@@ -17,6 +42,14 @@ BG_BLUE = 44
 BG_MAGENTA = 45
 BG_CYAN = 46
 BG_WHITE = 47
+BG_BRIGHT_BLACK = 100
+BG_BRIGHT_RED = 101
+BG_BRIGHT_GREEN = 102
+BG_BRIGHT_YELLOW = 103
+BG_BRIGHT_BLUE = 104
+BG_BRIGHT_MAGENTA = 105
+BG_BRIGHT_CYAN = 106
+BG_BRIGHT_WHITE = 107
 
 VALID_BACKGROUNDS = [
     BG_BLACK,
@@ -27,6 +60,14 @@ VALID_BACKGROUNDS = [
     BG_MAGENTA,
     BG_CYAN,
     BG_WHITE,
+    BG_BRIGHT_BLACK,
+    BG_BRIGHT_RED,
+    BG_BRIGHT_GREEN,
+    BG_BRIGHT_YELLOW,
+    BG_BRIGHT_BLUE,
+    BG_BRIGHT_MAGENTA,
+    BG_BRIGHT_CYAN,
+    BG_BRIGHT_WHITE,
 ]
 
 BOLD = 1  # Bold or increased intensity
@@ -101,6 +142,34 @@ def stylize_section_header(text: str):
     return stylize(f" {text} ", color=BLACK, background=BG_WHITE, style=UNDERLINE)
 
 
+def stylize_green(text: str):
+    return stylize(text, color=GREEN)
+
+
+def stylize_blue(text: str):
+    return stylize(text, color=BLUE)
+
+
+def stylize_cyan(text: str):
+    return stylize(text, color=CYAN)
+
+
+def stylize_magenta(text: str):
+    return stylize(text, color=MAGENTA)
+
+
+def stylize_yellow(text: str):
+    return stylize(text, color=YELLOW)
+
+
+def stylize_red(text: str):
+    return stylize(text, color=RED)
+
+
+def stylize_bold_green(text: str):
+    return stylize(text, color=GREEN, style=BOLD)
+
+
 def stylize_bold_yellow(text: str):
     return stylize(text, color=YELLOW, style=BOLD)
 

zrb/util/cli/subcommand.py

@@ -1,5 +1,5 @@
-from ...group.any_group import AnyGroup
-from ..group import get_non_empty_subgroups, get_subtasks
+from zrb.group.any_group import AnyGroup
+from zrb.util.group import get_non_empty_subgroups, get_subtasks
 
 
 class SubCommand:

zrb/util/codemod/add_code_to_class.py

@@ -0,0 +1,35 @@
+import libcst as cst
+
+
+class ClassCodeAdder(cst.CSTTransformer):
+    def __init__(self, class_name: str, new_code: str):
+        self.class_name = class_name
+        self.new_code = cst.parse_statement(new_code)
+        self.class_found = False
+
+    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
+            # Add the method to the class body
+            new_body = updated_node.body.with_changes(
+                body=updated_node.body.body + (self.new_code,)
+            )
+            return updated_node.with_changes(body=new_body)
+        return updated_node
+
+
+def add_code_to_class(original_code: str, class_name: str, method_code: str) -> str:
+    # Parse the original code into a module
+    module = cst.parse_module(original_code)
+    # Initialize transformer with the class name and method code
+    transformer = ClassCodeAdder(class_name, method_code)
+    # 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

zrb/util/codemod/add_code_to_function.py

@@ -0,0 +1,36 @@
+import libcst as cst
+
+
+class FunctionCodeAdder(cst.CSTTransformer):
+    def __init__(self, function_name: str, new_code: str):
+        self.function_name = function_name
+        # Use parse_module to handle multiple statements
+        self.new_code = cst.parse_statement(new_code)
+        self.function_found = False
+
+    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
+            # Add the method to the class body
+            new_body = updated_node.body.with_changes(
+                body=updated_node.body.body + (self.new_code,)
+            )
+            return updated_node.with_changes(body=new_body)
+        return updated_node
+
+
+def add_code_to_function(original_code: str, function_name: str, new_code: str) -> str:
+    # Parse the original code into a module
+    module = cst.parse_module(original_code)
+    # Initialize the transformer with the necessary information
+    transformer = FunctionCodeAdder(function_name, new_code)
+    # 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

zrb/util/codemod/add_code_to_method.py

@@ -0,0 +1,55 @@
+import libcst as cst
+
+
+class MethodCodeAdder(cst.CSTTransformer):
+    def __init__(self, class_name: str, method_name: str, new_code: str):
+        self.class_name = class_name
+        self.method_name = method_name
+        # Use parse_module to handle multiple statements
+        self.new_code = cst.parse_statement(new_code)
+        self.class_found = False
+        self.method_found = False
+
+    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
+                    body_with_new_code = item.body.with_changes(
+                        body=item.body.body + (self.new_code,)  # Add the 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
+
+
+def add_code_to_method(
+    original_code: str, class_name: str, function_name: str, new_code: str
+) -> str:
+    # Parse the original code into a module
+    module = cst.parse_module(original_code)
+    # Initialize the transformer with the necessary information
+    transformer = MethodCodeAdder(class_name, function_name, new_code)
+    # 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 {function_name} not found in class {class_name}.")
+    # Return the modified code
+    return modified_module.code

zrb/util/codemod/add_key_to_dict.py

@@ -0,0 +1,51 @@
+import libcst as cst
+
+
+class DictionaryModifier(cst.CSTTransformer):
+    def __init__(self, dictionary_name: str, new_key: str, new_value: str):
+        self.dictionary_name = dictionary_name
+        self.new_key = new_key
+        self.new_value = new_value
+        self.found = False
+
+    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
+                # Add new key-value pair to the existing dictionary
+                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
+
+
+def add_key_to_dict(
+    original_code: str, dictionary_name: str, new_key: str, new_value: str
+) -> 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)
+    # 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

zrb/util/codemod/add_param_to_function_call.py

@@ -0,0 +1,39 @@
+import libcst as cst
+
+
+class FunctionCallParamAdder(cst.CSTTransformer):
+    def __init__(self, func_name: str, new_param: str):
+        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
+
+    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
+        ):  # noqa
+            # Add the new parameter to the function call arguments
+            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
+
+
+def add_param_to_function_call(
+    original_code: str, func_name: str, new_param: str
+) -> str:
+    # Parse the original code into a module
+    module = cst.parse_module(original_code)
+    # Initialize the transformer with the necessary information
+    transformer = FunctionCallParamAdder(func_name, new_param)
+    # 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 {func_name} not found in the provided code."
+        )
+    # Return the modified code
+    return modified_module.code

zrb/util/codemod/add_property_to_class.py

@@ -0,0 +1,55 @@
+import libcst as cst
+
+
+class ClassPropertyAdder(cst.CSTTransformer):
+    def __init__(
+        self, class_name: str, property_name: str, annotation: str, default_value: str
+    ):
+        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
+
+    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,
+                    )
+                ]
+            )
+            # Insert the new property at the start of the class body with a newline
+            new_body = cst.IndentedBlock(body=(new_property,) + updated_node.body.body)
+            return updated_node.with_changes(body=new_body)
+        return updated_node
+
+
+def add_property_to_class(
+    original_code: str,
+    class_name: str,
+    property_name: str,
+    annotation: str,
+    default_value: str,
+) -> 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 = ClassPropertyAdder(
+        class_name, property_name, annotation, default_value
+    )
+    # 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

zrb/util/git.py

@@ -0,0 +1,156 @@
+import subprocess
+
+from pydantic import BaseModel
+
+
+class DiffResult(BaseModel):
+    created: list[str]
+    removed: list[str]
+    updated: list[str]
+
+
+def get_diff(source_commit: str, current_commit: str) -> DiffResult:
+    # git show b176b5a main
+    exit_status, output = subprocess.getstatusoutput(
+        f"git diff {source_commit} {current_commit}"
+    )
+    if exit_status != 0:
+        raise Exception(output)
+    lines = output.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"]
+        ],
+    )
+
+
+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 result.stdout.strip()
+    except subprocess.CalledProcessError as e:
+        raise Exception(e.stderr or e.stdout)
+
+
+def get_current_branch() -> str:
+    try:
+        result = subprocess.run(
+            ["git", "rev-parse", "--abbrev-ref", "HEAD"],
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            text=True,
+            check=True,
+        )
+        return result.stdout.strip()
+    except subprocess.CalledProcessError as e:
+        raise Exception(e.stderr or e.stdout)
+
+
+def get_branches() -> list[str]:
+    try:
+        result = subprocess.run(
+            ["git", "branch"],
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            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(branch_name: str) -> str:
+    try:
+        result = subprocess.run(
+            ["git", "branch", "-D", branch_name],
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            text=True,
+            check=True,
+        )
+        return result.stdout.strip()
+    except subprocess.CalledProcessError as e:
+        raise Exception(e.stderr or e.stdout)
+
+
+def add() -> str:
+    try:
+        subprocess.run(
+            ["git", "add", ".", "-A"],
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            text=True,
+            check=True,
+        )
+    except subprocess.CalledProcessError as e:
+        raise Exception(e.stderr or e.stdout)
+
+
+def commit(message: str) -> str:
+    try:
+        subprocess.run(
+            ["git", "commit", "-m", message],
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            text=True,
+            check=True,
+        )
+    except subprocess.CalledProcessError as e:
+        raise Exception(e.stderr or e.stdout)
+
+
+def pull(remote: str, branch: str) -> str:
+    try:
+        subprocess.run(
+            ["git", "pull", remote, branch],
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            text=True,
+            check=True,
+        )
+    except subprocess.CalledProcessError as e:
+        raise Exception(e.stderr or e.stdout)
+
+
+def push(remote: str, branch: str) -> str:
+    try:
+        subprocess.run(
+            ["git", "push", "-u", remote, branch],
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            text=True,
+            check=True,
+        )
+    except subprocess.CalledProcessError as e:
+        raise Exception(e.stderr or e.stdout)