computation-graph 52__tar.gz → 53__tar.gz

{computation-graph-52 → computation-graph-53}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: computation-graph
-Version: 52
+Version: 53
 Requires-Python: >=3
 Description-Content-Type: text/markdown
 License-File: LICENSE.md

{computation-graph-52 → computation-graph-53}/computation_graph/composers/duplication.py

@@ -1,6 +1,7 @@
 import asyncio
 import functools
 import inspect
+from typing import Dict
 
 import gamla
 
@@ -20,6 +21,7 @@ def duplicate_function(func):
     def inner(*args, **kwargs):
         return func(*args, **kwargs)
 
+    inner.__name__ = f"duplicate of {inner.__name__}"
     return inner
 
 
@@ -42,8 +44,8 @@ _duplicate_node = gamla.compose_left(
     graph.make_computation_node,
 )
 
+
 _node_to_duplicated_node = gamla.compose_left(
-    graph.get_all_nodes,
     gamla.remove(base_types.node_is_terminal),
     gamla.map(gamla.pair_right(_duplicate_node)),
     dict,
@@ -51,7 +53,7 @@ _node_to_duplicated_node = gamla.compose_left(
 )
 
 duplicate_graph = gamla.compose_left(
-    gamla.pair_with(_node_to_duplicated_node),
+    gamla.pair_with(gamla.compose_left(graph.get_all_nodes, _node_to_duplicated_node)),
     gamla.star(
         lambda get_duplicated_node, graph: gamla.pipe(
             graph,
@@ -68,3 +70,38 @@ duplicate_graph = gamla.compose_left(
 duplicate_function_or_graph = gamla.ternary(
     gamla.is_instance(tuple), duplicate_graph, duplicate_function
 )
+
+
+def safe_replace_sources(
+    source_to_replacement_dict: Dict[
+        base_types.CallableOrNode, base_types.CallableOrNodeOrGraph
+    ],
+    cg: base_types.GraphType,
+) -> base_types.GraphType:
+    source_to_replacement_dict = gamla.keymap(graph.make_computation_node)(
+        source_to_replacement_dict
+    )
+    reachable_to_duplicate_map = gamla.pipe(
+        gamla.graph_traverse_many(
+            source_to_replacement_dict.keys(), graph.traverse_forward(cg)
+        ),
+        gamla.remove(gamla.contains(source_to_replacement_dict.keys())),
+        _node_to_duplicated_node,
+    )
+    get_node_replacement = gamla.compose_left(
+        gamla.lazyjuxt(
+            gamla.dict_to_getter_with_default(None)(source_to_replacement_dict),
+            reachable_to_duplicate_map,
+            gamla.identity,
+        ),
+        gamla.find(gamla.identity),
+    )
+
+    def update_edge(e: base_types.ComputationEdge) -> base_types.GraphType:
+        dest = base_types.edge_destination(e)
+        g = graph.replace_node(dest, get_node_replacement(dest))((e,))
+        for source in base_types.edge_sources(e):
+            g = graph.replace_source(source, get_node_replacement(source), g)
+        return g
+
+    return gamla.reduce(lambda updated_cg, e: updated_cg + update_edge(e), (), cg)

{computation-graph-52 → computation-graph-53}/computation_graph/graph.py

@@ -1,5 +1,5 @@
 import dataclasses
-from typing import Callable, FrozenSet
+from typing import Callable, FrozenSet, Tuple
 
 import gamla
 from gamla.optimized import sync as opt_gamla
@@ -148,6 +148,25 @@ def _replace_source_in_edges(
     )
 
 
+traverse_forward: Callable[
+    [base_types.GraphType],
+    Callable[[base_types.CallableOrNode], Tuple[base_types.ComputationNode, ...]],
+] = gamla.compose_left(
+    gamla.mapcat(
+        gamla.compose_left(
+            gamla.juxt(base_types.edge_sources, base_types.edge_destination),
+            gamla.explode(0),
+        )
+    ),
+    gamla.groupby_many_reduce(
+        gamla.compose_left(gamla.head, gamla.wrap_tuple),
+        lambda destinations, e: (*(destinations if destinations else ()), e[1]),
+    ),
+    gamla.dict_to_getter_with_default(()),
+    gamla.before(make_computation_node),
+)
+
+
 @gamla.curry
 def replace_source(
     original: base_types.CallableOrNode,

{computation-graph-52 → computation-graph-53}/computation_graph/graph_test.py

@@ -1,4 +1,5 @@
 import asyncio
+from typing import Dict
 
 import gamla
 import pytest
@@ -802,3 +803,100 @@ def test_ambig_edges_assertion_in_merge_graphs_active_only_when_env_var_is_activ
             composers.compose_left_unary(lambda: 1, a),
             composers.compose_left_unary(lambda: 1, a),
         )
+
+
+def a():
+    pass
+
+
+def b(x):
+    pass
+
+
+def c(x):
+    pass
+
+
+def d(x, y):
+    pass
+
+
+def kuky():
+    pass
+
+
+def kuku():
+    pass
+
+
+t = graph.make_terminal("t", lambda x: x)
+
+g = (
+    composers.compose_left(a, c)
+    + composers.compose_left(c, d, key="x")
+    + composers.compose_left(b, d, key="y")
+    + composers.compose_left_future(d, b, "x", "bla")
+    + composers.compose_left(a, t)
+)
+
+
+@pytest.mark.parametrize(
+    "to_replace,expected_edges_strs",
+    [
+        pytest.param(
+            {a: kuky},
+            {
+                "kuky----x---->duplicate of c",
+                "kuky----x---->t",
+                "duplicate of c----x---->duplicate of d",
+                "when_memory_unavailable----x---->duplicate of b",
+                "duplicate of d....x....>duplicate of b",
+                "duplicate of b----y---->duplicate of d",
+            },
+            id="replace source node",
+        ),
+        pytest.param(
+            {c: kuky},
+            {
+                "a----x---->kuky",
+                "a----x---->t",
+                "kuky----x---->duplicate of d",
+                "when_memory_unavailable----x---->duplicate of b",
+                "duplicate of d....x....>duplicate of b",
+                "duplicate of b----y---->duplicate of d",
+            },
+            id="replace node not in cycle",
+        ),
+        pytest.param(
+            {b: kuky},
+            {
+                "a----x---->c",
+                "a----x---->t",
+                "c----x---->duplicate of d",
+                "when_memory_unavailable----x---->kuky",
+                "duplicate of d....x....>kuky",
+                "kuky----y---->duplicate of d",
+            },
+            id="replace node in cycle",
+        ),
+        pytest.param(
+            {a: kuku, b: kuky},
+            {
+                "duplicate of c----x---->duplicate of d",
+                "duplicate of d....x....>kuky",
+                "kuku----x---->duplicate of c",
+                "kuku----x---->t",
+                "kuky----y---->duplicate of d",
+                "when_memory_unavailable----x---->kuky",
+            },
+            id="replace multiple nodes - duplicate reachables once",
+        ),
+    ],
+)
+def test_safe_replace_node(
+    to_replace: Dict[base_types.CallableOrNode, base_types.CallableOrNodeOrGraph],
+    expected_edges_strs: str,
+):
+    assert expected_edges_strs == {
+        str(e) for e in duplication.safe_replace_sources(to_replace, g)
+    }

{computation-graph-52 → computation-graph-53}/computation_graph.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: computation-graph
-Version: 52
+Version: 53
 Requires-Python: >=3
 Description-Content-Type: text/markdown
 License-File: LICENSE.md

{computation-graph-52 → computation-graph-53}/setup.py

@@ -7,7 +7,7 @@ with open("README.md", "r") as fh:
 setuptools.setup(
     name="computation-graph",
     python_requires=">=3",
-    version="52",
+    version="53",
     long_description=_LONG_DESCRIPTION,
     long_description_content_type="text/markdown",
     packages=setuptools.find_namespace_packages(),