zrb 0.6.0__py3-none-any.whl → 0.6.1__py3-none-any.whl

zrb/helper/codemod/add_import_module.py

@@ -1,14 +1,15 @@
-from zrb.helper.typecheck import typechecked
-from zrb.helper.typing import Optional, Union, Tuple
 import libcst as cst
 
+from zrb.helper.typecheck import typechecked
+from zrb.helper.typing import Optional, Tuple, Union
+
 
 @typechecked
 def add_import_module(
     code: str,
     module_path: str,
     resource: Optional[str] = None,
-    alias: Optional[str] = None
+    alias: Optional[str] = None,
 ) -> str:
     """
     Parses the given code as a module using `libcst.parse_module()`,
@@ -29,16 +30,14 @@ def add_import_module(
     """
     module = cst.parse_module(code)
     new_import = _get_new_import(
-        module_path=module_path,
-        resource=resource,
-        alias=alias
+        module_path=module_path, resource=resource, alias=alias
     )
 
     last_import_index = None
     for i, node in enumerate(module.body):
         if isinstance(node, cst.SimpleStatementLine) and (
-            isinstance(node.body[0], cst.Import) or
-            isinstance(node.body[0], cst.ImportFrom)
+            isinstance(node.body[0], cst.Import)
+            or isinstance(node.body[0], cst.ImportFrom)
         ):
             last_import_index = i
 
@@ -58,44 +57,34 @@ def add_import_module(
 
 
 def _get_new_import(
-    module_path: str,
-    resource: Optional[str] = None,
-    alias: Optional[str] = None
+    module_path: str, resource: Optional[str] = None, alias: Optional[str] = None
 ) -> Union[cst.ImportFrom, cst.Import]:
     dots, module_path = _split_module_path(module_path)
     relative = [cst.Dot()] * len(dots)
     if resource is None:
         if len(relative) > 0:
-            raise Exception(
-                'Relative import is not allowed, please specify resource'
-            )
+            raise Exception("Relative import is not allowed, please specify resource")
         return cst.Import(
             names=[
                 cst.ImportAlias(
-                    name=cst.parse_expression(module_path),
-                    asname=_get_as_name(alias)
+                    name=cst.parse_expression(module_path), asname=_get_as_name(alias)
                 )
             ],
         )
     return cst.ImportFrom(
         module=cst.parse_expression(module_path),
-        names=[
-            cst.ImportAlias(
-                name=cst.Name(resource),
-                asname=_get_as_name(alias)
-            )
-        ],
-        relative=relative
+        names=[cst.ImportAlias(name=cst.Name(resource), asname=_get_as_name(alias))],
+        relative=relative,
     )
 
 
 @typechecked
 def _split_module_path(module_path) -> Tuple[str, str]:
-    prefix = ''
-    suffix = ''
+    prefix = ""
+    suffix = ""
     is_prefix = True
     for char in module_path:
-        if char != '.':
+        if char != ".":
             is_prefix = False
         if is_prefix:
             prefix += char

zrb/helper/codemod/add_key_value_to_dict.py

@@ -1,13 +1,15 @@
-from zrb.helper.typecheck import typechecked
 import libcst as cst
 
+from zrb.helper.typecheck import typechecked
+
 
 @typechecked
 class AddKeyValuePairTransformer(cst.CSTTransformer):
     def __init__(
-        self, dict_name: str,
+        self,
+        dict_name: str,
         key_node: cst.BaseExpression,
-        value_node: cst.BaseExpression
+        value_node: cst.BaseExpression,
     ):
         self.dict_name = dict_name
         self.key_node = key_node
@@ -24,10 +26,7 @@ class AddKeyValuePairTransformer(cst.CSTTransformer):
         ):
             for element in updated_node.value.elements:
                 key = element.key
-                if (
-                    isinstance(key, cst.BaseExpression)
-                    and key.value == self.key_node
-                ):
+                if isinstance(key, cst.BaseExpression) and key.value == self.key_node:
                     return updated_node
             new_elements = updated_node.value.elements + (
                 cst.DictElement(self.key_node, self.value_node),
@@ -37,9 +36,7 @@ class AddKeyValuePairTransformer(cst.CSTTransformer):
         return updated_node
 
 
-def add_key_value_to_dict(
-    code: str, dict_name: str, key: str, value: str
-) -> str:
+def add_key_value_to_dict(code: str, dict_name: str, key: str, value: str) -> str:
     module = cst.parse_module(code)
     key_node = cst.parse_expression(key)
     value_node = cst.parse_expression(value)

zrb/helper/codemod/add_property_to_class.py

@@ -1,6 +1,7 @@
+import libcst as cst
+
 from zrb.helper.typecheck import typechecked
 from zrb.helper.typing import Optional
-import libcst as cst
 
 
 @typechecked
@@ -10,7 +11,7 @@ class AddPropertyTransformer(cst.CSTTransformer):
         class_name: str,
         property_name_node: cst.Name,
         property_type_node: cst.Annotation,
-        property_value_node: Optional[cst.BaseExpression]
+        property_value_node: Optional[cst.BaseExpression],
     ):
         self.class_name = class_name
         self.property_name_node = property_name_node
@@ -27,9 +28,10 @@ class AddPropertyTransformer(cst.CSTTransformer):
         if not isinstance(original_node.body, cst.EmptyLine):
             for stmt in original_node.body.body:
                 if (
-                    hasattr(stmt, 'body') and
-                    len(stmt.body) > 0 and
-                    isinstance(stmt.body[0], cst.Pass)
+                    hasattr(stmt, "body")
+                    and not isinstance(stmt.body, cst.IndentedBlock)
+                    and len(stmt.body) > 0
+                    and isinstance(stmt.body[0], cst.Pass)
                 ):
                     continue
                 # Add all old class body except any `pass` statement
@@ -59,7 +61,7 @@ def add_property_to_class(
     class_name: str,
     property_name: str,
     property_type: str,
-    property_value: Optional[str] = None
+    property_value: Optional[str] = None,
 ) -> str:
     module = cst.parse_module(code)
     property_name_node = cst.Name(value=property_name)
@@ -70,14 +72,14 @@ def add_property_to_class(
             class_name=class_name,
             property_name_node=property_name_node,
             property_type_node=property_type_node,
-            property_value_node=property_value_node
+            property_value_node=property_value_node,
         )
     )
     return transformed_module.code
 
 
 def _get_property_value_node(
-    property_value: Optional[str]
+    property_value: Optional[str],
 ) -> Optional[cst.BaseExpression]:
     if property_value is None:
         return None

zrb/helper/codemod/add_upstream_to_task.py

@@ -1,6 +1,7 @@
-from zrb.helper.typecheck import typechecked
 import libcst as cst
 
+from zrb.helper.typecheck import typechecked
+
 
 @typechecked
 class TaskTransformer(cst.CSTTransformer):
@@ -12,10 +13,7 @@ class TaskTransformer(cst.CSTTransformer):
         self._upstream_added = False
 
     def _is_task(self, node: cst.Call) -> bool:
-        if (
-            not isinstance(node.func, cst.Name)
-            or node.func.value != "Task"
-        ):
+        if not isinstance(node.func, cst.Name) or node.func.value != "Task":
             return False
         # Look for name keyword
         for name_arg in node.args:
@@ -33,9 +31,7 @@ class TaskTransformer(cst.CSTTransformer):
         if self._is_task(node):
             self._on_task_call = True
 
-    def leave_Call(
-        self, original_node: cst.Call, updated_node: cst.Call
-    ) -> cst.Call:
+    def leave_Call(self, original_node: cst.Call, updated_node: cst.Call) -> cst.Call:
         # Remove flags.
         # If upstreams argument not found (see: self.leave_Arg), add it
         if not self._is_task(updated_node):
@@ -44,16 +40,16 @@ class TaskTransformer(cst.CSTTransformer):
         if self._upstream_added:
             return updated_node
         new_args = [arg for arg in updated_node.args]
-        new_args.append(cst.Arg(
-            keyword=cst.parse_expression('upstreams'),
-            value=cst.List([
-                cst.Element(cst.parse_expression(self.upstream_task_name))
-            ])
-        ))
-        self._upstream_added = True
-        return updated_node.with_changes(
-            args=new_args
+        new_args.append(
+            cst.Arg(
+                keyword=cst.parse_expression("upstreams"),
+                value=cst.List(
+                    [cst.Element(cst.parse_expression(self.upstream_task_name))]
+                ),
+            )
         )
+        self._upstream_added = True
+        return updated_node.with_changes(args=new_args)
 
     def leave_Arg(self, orginnal_node: cst.Arg, updated_node: cst.Arg):
         # Once the flag was set (see: self.visit_Call), look for
@@ -61,25 +57,19 @@ class TaskTransformer(cst.CSTTransformer):
         if (
             not self._on_task_call
             or updated_node.keyword is None
-            or updated_node.keyword.value != 'upstreams'
+            or updated_node.keyword.value != "upstreams"
         ):
             return updated_node
-        new_elements = [
-            old_element for old_element in updated_node.value.elements
-        ]
-        new_elements.append(cst.Element(
-            value=cst.parse_expression(self.upstream_task_name)
-        ))
-        self._upstream_added = True
-        return updated_node.with_changes(
-            value=cst.List(new_elements)
+        new_elements = [old_element for old_element in updated_node.value.elements]
+        new_elements.append(
+            cst.Element(value=cst.parse_expression(self.upstream_task_name))
         )
+        self._upstream_added = True
+        return updated_node.with_changes(value=cst.List(new_elements))
 
 
 @typechecked
-def add_upstream_to_task(
-    code: str, task_name: str, upstream_task_name: str
-) -> str:
+def add_upstream_to_task(code: str, task_name: str, upstream_task_name: str) -> str:
     """
     Transformer to add an upstream task to a Task object.
 
@@ -91,9 +81,7 @@ def add_upstream_to_task(
     Returns:
         A modified version of the input code with the upstream task added.
     """
-    transformer: TaskTransformer = TaskTransformer(
-        task_name, upstream_task_name
-    )
+    transformer: TaskTransformer = TaskTransformer(task_name, upstream_task_name)
     original_tree: cst.Module = cst.parse_module(code)
     modified_tree: cst.Module = original_tree.visit(transformer)
     return modified_tree.code

zrb/helper/codemod/append_code_to_function.py

@@ -1,6 +1,7 @@
-from zrb.helper.typecheck import typechecked
 import libcst as cst
 
+from zrb.helper.typecheck import typechecked
+
 
 @typechecked
 class AppendCodeTransformer(cst.CSTTransformer):
@@ -13,21 +14,13 @@ class AppendCodeTransformer(cst.CSTTransformer):
         self, original_node: cst.FunctionDef, updated_node: cst.FunctionDef
     ) -> cst.FunctionDef:
         if updated_node.name.value == self.function_name:
-            new_body = updated_node.body.body + (
-                cst.parse_statement(self.new_code),
-            )
-            return updated_node.with_changes(
-                body=cst.IndentedBlock(body=new_body)
-            )
+            new_body = updated_node.body.body + (cst.parse_statement(self.new_code),)
+            return updated_node.with_changes(body=cst.IndentedBlock(body=new_body))
         return updated_node
 
 
 @typechecked
-def append_code_to_function(
-    code: str, function_name: str, new_code: str
-) -> str:
+def append_code_to_function(code: str, function_name: str, new_code: str) -> str:
     module = cst.parse_module(code)
-    transformed_module = module.visit(
-        AppendCodeTransformer(function_name, new_code)
-    )
+    transformed_module = module.visit(AppendCodeTransformer(function_name, new_code))
     return transformed_module.code

zrb/helper/codemod/format_code.py

@@ -1,6 +1,7 @@
-from zrb.helper.typecheck import typechecked
 import autopep8
 
+from zrb.helper.typecheck import typechecked
+
 
 @typechecked
 def format_code(code: str) -> str:

zrb/helper/default_env.py

@@ -1,58 +1,58 @@
+import logging
+import os
+from functools import lru_cache
+
 from zrb.config.config import logging_level
-from zrb.helper.typecheck import typechecked
-from zrb.helper.typing import Mapping, Optional
 from zrb.helper.accessories.color import colored
 from zrb.helper.log import logger
-from functools import lru_cache
-
-import logging
-import os
+from zrb.helper.typecheck import typechecked
+from zrb.helper.typing import Mapping, Optional
 
 _PROJECT_DIR_MAP: Mapping[str, str] = {}
 
 
 def inject_default_env():
     # Inject PYTHONUNBUFFERED
-    if 'PYTHONUNBUFFERED' not in os.environ:
+    if "PYTHONUNBUFFERED" not in os.environ:
         if logging_level <= logging.INFO:
-            logger.info(colored('Set PYTHONUNBUFFERED to 1', attrs=['dark']))
-        os.environ['PYTHONUNBUFFERED'] = '1'
+            logger.info(colored("Set PYTHONUNBUFFERED to 1", attrs=["dark"]))
+        os.environ["PYTHONUNBUFFERED"] = "1"
     # Inject ZRB_HOME_DIR
-    if 'ZRB_HOME_DIR' not in os.environ:
+    if "ZRB_HOME_DIR" not in os.environ:
         default_home_dir = os.path.dirname(os.path.dirname(__file__))
         if logging_level <= logging.INFO:
-            logger.info(colored(
-                f'Set ZRB_HOME_DIR to {default_home_dir}', attrs=['dark']
-            ))
-        os.environ['ZRB_HOME_DIR'] = default_home_dir
+            logger.info(
+                colored(f"Set ZRB_HOME_DIR to {default_home_dir}", attrs=["dark"])
+            )
+        os.environ["ZRB_HOME_DIR"] = default_home_dir
     # Inject ZRB_PROJECT_DIR
     current_dir = os.getcwd()
     if current_dir not in _PROJECT_DIR_MAP:
         if logging_level <= logging.INFO:
-            logger.info(colored('Getting project directory', attrs=['dark']))
+            logger.info(colored("Getting project directory", attrs=["dark"]))
         zrb_project_dir = _get_project_dir(current_dir)
         if zrb_project_dir is None:
             zrb_project_dir = current_dir
         _PROJECT_DIR_MAP[current_dir] = zrb_project_dir
         if logging_level <= logging.INFO:
-            logger.info(colored(
-                f'Set ZRB_PROJECT_DIR to {zrb_project_dir}', attrs=['dark']
-            ))
-        os.environ['ZRB_PROJECT_DIR'] = zrb_project_dir
+            logger.info(
+                colored(f"Set ZRB_PROJECT_DIR to {zrb_project_dir}", attrs=["dark"])
+            )
+        os.environ["ZRB_PROJECT_DIR"] = zrb_project_dir
     # Inject ZRB_PROJECT_NAME
-    if 'ZRB_PROJECT_NAME' not in os.environ:
+    if "ZRB_PROJECT_NAME" not in os.environ:
         zrb_project_name = os.path.basename(zrb_project_dir)
         if logging_level <= logging.INFO:
-            logger.info(colored(
-                f'Set ZRB_PROJECT_NAME to {zrb_project_name}', attrs=['dark']
-            ))
-        os.environ['ZRB_PROJECT_NAME'] = zrb_project_name
+            logger.info(
+                colored(f"Set ZRB_PROJECT_NAME to {zrb_project_name}", attrs=["dark"])
+            )
+        os.environ["ZRB_PROJECT_NAME"] = zrb_project_name
 
 
 @lru_cache
 @typechecked
 def _get_project_dir(project_dir: str) -> Optional[str]:
-    project_script = os.path.join(project_dir, 'zrb_init.py')
+    project_script = os.path.join(project_dir, "zrb_init.py")
     if os.path.isfile(project_script):
         return project_dir
     # zrb_init.py not found, look for it on the parent directory

zrb/helper/docker_compose/fetch_external_env.py

@@ -5,15 +5,15 @@ from zrb.helper.typing import Any, List, Mapping
 @typechecked
 def fetch_compose_file_env_map(data: Any) -> Mapping[str, str]:
     global_env_dict = {}
-    if 'services' not in data:
+    if "services" not in data:
         return global_env_dict
-    for service in data['services']:
-        if 'environment' not in data['services'][service]:
+    for service in data["services"]:
+        if "environment" not in data["services"][service]:
             continue
-        environments = data['services'][service]['environment']
+        environments = data["services"][service]["environment"]
         if isinstance(environments, list):
             for environment in environments:
-                parts: List[str] = environment.split('=')
+                parts: List[str] = environment.split("=")
                 if len(parts) > 0:
                     env_str = str(parts[1])
                     env_dict = parse_compose_file_env_string(env_str)
@@ -30,29 +30,29 @@ def fetch_compose_file_env_map(data: Any) -> Mapping[str, str]:
 def parse_compose_file_env_string(env_str: str) -> Mapping[str, str]:
     env_dict = {}
     stack = []
-    key = ''
-    value = ''
+    key = ""
+    value = ""
     index = 0
     while index < len(env_str):
         char = env_str[index]
-        if char == '{':
+        if char == "{":
             stack.append(index)
-        elif char == '}':
+        elif char == "}":
             start = stack.pop()
             if not stack:
-                segment = env_str[start+1:index]
-                if ':-' in segment:
-                    key, value = segment.split(':-', 1)
-                    if value.startswith('${') and value.endswith('}'):
+                segment = env_str[start + 1 : index]
+                if ":-" in segment:
+                    key, value = segment.split(":-", 1)
+                    if value.startswith("${") and value.endswith("}"):
                         sub_dict = parse_compose_file_env_string(value)
                         env_dict.update(sub_dict)
                 else:
                     key = segment
-                    value = ''
-                if '${' in value:
-                    value = ''
+                    value = ""
+                if "${" in value:
+                    value = ""
                 env_dict[key] = value
-                key = ''
-                value = ''
+                key = ""
+                value = ""
         index += 1
     return env_dict

zrb/helper/docker_compose/file.py

@@ -1,12 +1,13 @@
+from ruamel.yaml import YAML, CommentedMap
+
 from zrb.helper.typecheck import typechecked
 from zrb.helper.typing import Any, Mapping
-from ruamel.yaml import YAML, CommentedMap
 
 
 @typechecked
 def read_compose_file(file_name: str) -> Any:
     yaml = YAML()
-    with open(file_name, 'r') as f:
+    with open(file_name, "r") as f:
         data = yaml.load(f)
     return data
 
@@ -14,7 +15,7 @@ def read_compose_file(file_name: str) -> Any:
 @typechecked
 def write_compose_file(file_name: str, data: Any):
     yaml = YAML()
-    with open(file_name, 'w') as f:
+    with open(file_name, "w") as f:
         yaml.dump(data, f)
 
 
@@ -22,5 +23,5 @@ def write_compose_file(file_name: str, data: Any):
 def add_services(file_name: str, new_services: Mapping[str, str]):
     data = read_compose_file(file_name)
     data = CommentedMap(data)
-    data['services'].update(CommentedMap(new_services))
+    data["services"].update(CommentedMap(new_services))
     write_compose_file(file_name, data)