feat: add envs_depth to procs to control the depth of envs to be in…

…herited by subclasses

pipen/proc.py

@@ -105,6 +105,7 @@ class Proc(ABC, metaclass=ProcMeta):
             the docstring by default.
         envs: The arguments that are job-independent, useful for common options
             across jobs.
+        envs_depth: How deep to update the envs when subclassed.
         cache: Should we detect whether the jobs are cached?
         dirsig: When checking the signature for caching, whether should we walk
             through the content of the directory? This is sometimes
@@ -153,6 +154,7 @@ class Proc(ABC, metaclass=ProcMeta):
     name: str = None
     desc: str = None
     envs: Mapping[str, Any] = None
+    envs_depth: int = None
     cache: bool = None
     dirsig: bool = None
     export: bool = None
@@ -188,6 +190,7 @@ def from_proc(
         name: str = None,
         desc: str = None,
         envs: Mapping[str, Any] = None,
+        envs_depth: int = None,
         cache: bool = None,
         export: bool = None,
         error_strategy: str = None,
@@ -209,6 +212,7 @@ def from_proc(
             desc: The new description of the process
             envs: The arguments of the process, will overwrite parent one
                 The items that are specified will be inherited
+            envs_depth: How deep to update the envs when subclassed.
             cache: Whether we should check the cache for the jobs
             export: When True, the results will be exported to
                 `<pipeline.outdir>`
@@ -257,6 +261,7 @@ def from_proc(
         for key in (
             "desc",
             "envs",
+            "envs_depth",
             "cache",
             "forks",
             "order",
@@ -298,7 +303,12 @@ def __init_subclass__(cls) -> None:
                 r"'^[\w.-]+$'"
             )
 
-        envs = update_dict(parent.envs if parent else None, cls.envs)
+        envs = update_dict(
+            parent.envs if parent else None,
+            cls.envs,
+            depth=0 if not parent or parent.envs_depth is None else parent.envs_depth,
+        )
+        # So values can be accessed like Proc.envs.a.b
         cls.envs = envs if isinstance(envs, Diot) else Diot(envs or {})
         cls.plugin_opts = update_dict(
             parent.plugin_opts if parent else None,

pipen/utils.py

@@ -206,13 +206,14 @@ def desc_from_docstring(
 def update_dict(
     parent: Mapping[str, Any],
     new: Mapping[str, Any],
+    depth: int = 0,
 ) -> Mapping[str, Any]:
     """Update the new dict to the parent, but make sure parent does not change
 
     Args:
         parent: The parent dictionary
         new: The new dictionary
-        depth: The depth to be copied.
+        depth: The depth to be copied. 0 for updating to the deepest level.
 
     Examples:
         >>> parent = {"a": {"b": 1}}
@@ -232,10 +233,11 @@ def update_dict(
             key not in out
             or not isinstance(val, dict)
             or not isinstance(out[key], dict)
+            or depth == 1
         ):
             out[key] = val
         else:
-            out[key] = update_dict(out[key], val)
+            out[key] = update_dict(out[key], val, depth - 1)
 
     return out
 

tests/test_proc.py

@@ -157,11 +157,34 @@ def test_invalid_name():
 
 
 def test_inherit_proc_envs():
-    class Proc1(Proc):
+    class Proc1_1(Proc):
+        envs_depth = 1
         envs = {"a": {"b": 1, "c": 2}}
 
-    class Proc2(Proc1):
+    class Proc2(Proc1_1):
         envs = {"a": {"b": 2}}
 
-    assert Proc2.envs == {"a": {"b": 2, "c": 2}}
-    assert Proc1.envs == {"a": {"b": 1, "c": 2}}
+    class Proc1_2(Proc):
+        envs_depth = 2
+        envs = {"a": {"b": 1, "c": 2}}
+
+    class Proc3(Proc1_2):
+        envs_depth = 2
+        envs = {"a": {"b": 3}}
+
+    class Proc1_3(Proc):
+        envs_depth = 3
+        envs = {"a": {"b": 1, "c": 2}}
+
+    class Proc4(Proc1_3):
+        envs = {"a": {"b": 4}}
+
+    Proc5 = Proc.from_proc(Proc1_3, envs={"a": {"b": 5}})
+
+    assert Proc5.envs == {"a": {"b": 5, "c": 2}}
+    assert Proc4.envs == {"a": {"b": 4, "c": 2}}
+    assert Proc3.envs == {"a": {"b": 3, "c": 2}}
+    assert Proc2.envs == {"a": {"b": 2}}
+    assert Proc1_1.envs == {"a": {"b": 1, "c": 2}}
+    assert Proc1_2.envs == {"a": {"b": 1, "c": 2}}
+    assert Proc1_3.envs == {"a": {"b": 1, "c": 2}}

tests/test_utils.py

@@ -80,6 +80,25 @@ def test_update_dict():
     assert update_dict({"a": {"b": 1}}, {"a": {"c": 2}}) == {
         "a": {"b": 1, "c": 2}
     }
+    assert update_dict({"a": {"b": 1}}, {"a": {"c": 2}}, depth=1) == {
+        "a": {"c": 2}
+    }
+    assert update_dict(
+        {"a": {"b1": {"c": 1, "d": 2}, "b2": {"c": 1, "d": 2}}},
+        {"a": {"b1": {"c": 2}}},
+    ) == {"a": {"b1": {"c": 2, "d": 2}, "b2": {"c": 1, "d": 2}}}
+
+    assert update_dict(
+        {"a": {"b1": {"c": 1, "d": 2}, "b2": {"c": 1, "d": 2}}},
+        {"a": {"b1": {"c": 2}}},
+        depth=2,
+    ) == {"a": {"b1": {"c": 2}, "b2": {"c": 1, "d": 2}}}
+
+    assert update_dict(
+        {"a": {"b1": {"c": 1, "d": 2}, "b2": {"c": 1, "d": 2}}},
+        {"a": {"b1": {"c": 2}}},
+        depth=1,
+    ) == {"a": {"b1": {"c": 2}}}
 
 
 def test_strsplit():
@@ -146,7 +165,7 @@ def test_get_obj_from_spec():
     obj = _get_obj_from_spec(f"{HERE}/helpers.py:SimpleProc")
     assert obj.name == "SimpleProc"
 
-    obj = _get_obj_from_spec(f"pipen:Pipen")
+    obj = _get_obj_from_spec("pipen:Pipen")
     assert obj is pipen.Pipen