zrb 0.26.2__py3-none-any.whl → 1.0.0a5__py3-none-any.whl

zrb/util/cli/style.py

@@ -0,0 +1,194 @@
+BLACK = 30
+RED = 31
+GREEN = 32
+YELLOW = 33
+BLUE = 34
+MAGENTA = 35
+CYAN = 36
+WHITE = 37
+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
+BG_GREEN = 42
+BG_YELLOW = 43
+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,
+    BG_RED,
+    BG_GREEN,
+    BG_YELLOW,
+    BG_BLUE,
+    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
+FAINT = 2  # Faint, decreased intensity, or dim
+ITALIC = 3  # Italicized (not widely supported)
+UNDERLINE = 4  # Underline
+BLINK_SLOW = 5  # Slow blink (less than 150 per minute)
+BLINK_FAST = 6  # Rapid blink (MS-DOS style, not widely supported)
+REVERSED = 7  # Image negative, swap foreground and background colors
+HIDE = 8  # Conceal (not widely supported)
+CROSSED_OUT = 9  # Crossed-out or strike-through text
+
+VALID_STYLES = [
+    BOLD,
+    FAINT,
+    ITALIC,
+    UNDERLINE,
+    BLINK_SLOW,
+    BLINK_FAST,
+    REVERSED,
+    HIDE,
+    CROSSED_OUT,
+]
+
+ICONS = [
+    "🐶",
+    "🐱",
+    "🐭",
+    "🐹",
+    "🦊",
+    "🐻",
+    "🐨",
+    "🐯",
+    "🦁",
+    "🐮",
+    "🐷",
+    "🍎",
+    "🍐",
+    "🍊",
+    "🍋",
+    "🍌",
+    "🍉",
+    "🍇",
+    "🍓",
+    "🍈",
+    "🍒",
+    "🍑",
+]
+
+
+def stylize(
+    text: str,
+    color: int | None = None,
+    background: int | None = None,
+    style: int | None = None,
+):
+    # Start constructing the ANSI escape code
+    code_parts = []
+    if style is not None and style in VALID_STYLES:
+        code_parts.append(str(style))
+    if color is not None and color in VALID_COLORS:
+        code_parts.append(str(color))
+    if background is not None and background in VALID_BACKGROUNDS:
+        code_parts.append(str(background))
+    # Join all parts with ';' and add the escape code ending
+    if len(code_parts) > 0:
+        return "\033[" + ";".join(code_parts) + "m" + text + "\033[0m"
+    return text
+
+
+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)
+
+
+def stylize_bold_red(text: str):
+    return stylize(text, color=RED, style=BOLD)
+
+
+def stylize_faint(text: str):
+    return stylize(text, style=FAINT)
+
+
+def stylize_log(text: str):
+    return stylize_faint(text)
+
+
+def stylize_warning(text: str):
+    return stylize_bold_yellow(text)
+
+
+def stylize_error(text: str):
+    return stylize_bold_red(text)

zrb/util/cli/subcommand.py

@@ -0,0 +1,30 @@
+from zrb.group.any_group import AnyGroup
+from zrb.util.group import get_non_empty_subgroups, get_subtasks
+
+
+class SubCommand:
+    def __init__(self, paths: list[str] = [], nexts: list[str] = []):
+        self.paths = paths
+        self.nexts = nexts
+
+    def __repr__(self):
+        return f"<{self.__class__.__name__} paths={self.paths} nexts={self.nexts}>"
+
+
+def get_group_subcommands(
+    group: AnyGroup, previous_path: str = [], subcommands: list[SubCommand] = []
+) -> list[SubCommand]:
+    nexts = []
+    for task_alias in get_subtasks(group):
+        nexts.append(task_alias)
+    for subgroup_alias, subgroup in get_non_empty_subgroups(group).items():
+        nexts.append(subgroup_alias)
+        # Recursively add subgroup
+        get_group_subcommands(
+            group=subgroup,
+            previous_path=previous_path + [group.name],
+            subcommands=subcommands,
+        )
+    if len(nexts) > 0:
+        subcommands.append(SubCommand(paths=previous_path + [group.name], nexts=nexts))
+    return subcommands

zrb/util/cmd/remote.py

@@ -0,0 +1,20 @@
+import shlex
+
+
+def get_remote_cmd_script(
+    cmd_script: str,
+    host: str = "",
+    port: int | str = 22,
+    user: str = "",
+    password: str = "",
+    use_password: bool = False,
+    ssh_key: str = "",
+) -> str:
+    quoted_script = shlex.quote(cmd_script)
+    if ssh_key != "" and use_password:
+        return f'sshpass -p "{password}" ssh -t -p "{port}" -i "{ssh_key}" "{user}@{host}" {quoted_script}'  # noqa
+    if ssh_key != "":
+        return f'ssh -t -p "{port}" -i "{ssh_key}" "{user}@{host}" {quoted_script}'
+    if use_password:
+        return f'sshpass -p "{password}" ssh -t -p "{port}" "{user}@{host}" {quoted_script}'  # noqa
+    return f'ssh -t -p "{port}" "{user}@{host}" {quoted_script}'

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/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
+    )