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

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

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/{prepend_code_to_module.py → modify_module.py}

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

zrb/util/file.py

@@ -15,10 +15,11 @@ def write_file(file_path: str, content: str | list[str]):
         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)
-    # Remove trailing newlines, but keep one if it exists
     content = content.rstrip("\n")
-    if content.endswith("\n"):
+    if should_add_eol:
         content += "\n"
     with open(file_path, "w") as f:
         f.write(content)

{zrb-1.0.0b9.dist-info → zrb-1.0.0b10.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: zrb
-Version: 1.0.0b9
+Version: 1.0.0b10
 Summary: Your Automation Powerhouse
 Home-page: https://github.com/state-alchemists/zrb
 License: AGPL-3.0-or-later