zrb 1.0.0b9__py3-none-any.whl → 1.1.0__py3-none-any.whl

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/util/load.py

@@ -12,13 +12,17 @@ pattern = re.compile("[^a-zA-Z0-9]")
 def load_zrb_init(dir_path: str | None = None) -> Any | None:
     if dir_path is None:
         dir_path = os.getcwd()
-    script_path = os.path.join(dir_path, "zrb_init.py")
-    if os.path.isfile(script_path):
-        return load_file(script_path, -1)
-    new_dir_path = os.path.dirname(dir_path)
-    if new_dir_path == dir_path:
-        return
-    return load_zrb_init(new_dir_path)
+    # get path list from current path to the absolute root
+    current_path = os.path.abspath(dir_path)
+    path_list = [current_path]
+    while current_path != os.path.dirname(current_path):  # Stop at root
+        current_path = os.path.dirname(current_path)
+        path_list.append(current_path)
+    # loop from root to current path to load zrb_init
+    for current_path in path_list[::-1]:
+        script_path = os.path.join(current_path, "zrb_init.py")
+        if os.path.isfile(script_path):
+            load_file(script_path)
 
 
 @lru_cache
@@ -29,6 +33,8 @@ def load_file(script_path: str, sys_path_index: int = 0) -> Any | None:
     # Append script dir path
     script_dir_path = os.path.dirname(script_path)
     if script_dir_path not in sys.path:
+        if sys_path_index == -1:
+            sys_path_index = len(sys.path)
         sys.path.insert(sys_path_index, script_dir_path)
     # Add script dir path to Python path
     os.environ["PYTHONPATH"] = _get_new_python_path(script_dir_path)

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: zrb
-Version: 1.0.0b9
+Version: 1.1.0
 Summary: Your Automation Powerhouse
 Home-page: https://github.com/state-alchemists/zrb
 License: AGPL-3.0-or-later
@@ -26,7 +26,7 @@ Requires-Dist: pdfplumber (>=0.11.4,<0.12.0) ; extra == "rag"
 Requires-Dist: psutil (>=6.1.1,<7.0.0)
 Requires-Dist: pydantic-ai (>=0.0.19,<0.0.20)
 Requires-Dist: python-dotenv (>=1.0.1,<2.0.0)
-Requires-Dist: python-jose[cryptography] (>=3.3.0,<4.0.0)
+Requires-Dist: python-jose[cryptography] (>=3.4.0,<4.0.0)
 Requires-Dist: requests (>=2.32.3,<3.0.0)
 Requires-Dist: ulid-py (>=1.1.0,<2.0.0)
 Project-URL: Documentation, https://github.com/state-alchemists/zrb