computation-graph 67__tar.gz → 68__tar.gz

{computation_graph-67 → computation_graph-68}/PKG-INFO

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

computation_graph-68/computation_graph/composers/duplication.py

@@ -0,0 +1,154 @@
+import asyncio
+import dataclasses
+import functools
+from typing import Dict
+
+import gamla
+
+from computation_graph import base_types, graph
+from computation_graph.composers import debug
+
+
+def duplicate_function(func):
+    if asyncio.iscoroutinefunction(func):
+
+        @functools.wraps(func)
+        async def inner(*args, **kwargs):
+            return await func(*args, **kwargs)
+
+    else:
+
+        @functools.wraps(func)
+        def inner(*args, **kwargs):
+            return func(*args, **kwargs)
+
+    return debug.name_callable(inner, f"duplicate of {func.__name__}")
+
+
+def _duplicate_computation_edge(get_duplicated_node):
+    return gamla.compose_left(
+        gamla.dataclass_transform("source", get_duplicated_node),
+        gamla.dataclass_transform("destination", get_duplicated_node),
+        gamla.dataclass_transform(
+            "args", gamla.compose_left(gamla.map(get_duplicated_node), tuple)
+        ),
+    )
+
+
+def _signature_is_empty(signature: base_types.NodeSignature) -> bool:
+    return not (
+        signature.kwargs
+        or signature.optional_kwargs
+        or signature.is_args
+        or signature.is_kwargs
+    )
+
+
+def _duplicate_node(node: base_types.ComputationNode) -> base_types.ComputationNode:
+    if _signature_is_empty(node.signature):
+        return node
+    inner = duplicate_function(node.func)
+    return dataclasses.replace(node, name=inner.__name__, func=inner)
+
+
+def _node_to_duplicated_node(nodes):
+    return {node: _duplicate_node(node) for node in nodes if not node.is_terminal}.get
+
+
+def duplicate_graph(original: base_types.GraphType) -> base_types.GraphType:
+
+    _node_to_duplicated_node_index = gamla.pipe(
+        original.edges, graph.get_all_nodes, _node_to_duplicated_node
+    )
+    get_duplicated_node = gamla.when(
+        _node_to_duplicated_node_index, _node_to_duplicated_node_index
+    )
+    return base_types.GraphType(
+        edges=gamla.pipe(
+            original.edges,
+            gamla.map(
+                _duplicate_computation_edge(
+                    gamla.when(get_duplicated_node, get_duplicated_node)
+                )
+            ),
+            frozenset,
+        ),
+        sink=get_duplicated_node(original.sink),
+    )
+
+
+duplicate_function_or_graph = gamla.ternary(
+    gamla.is_instance(base_types.GraphType), 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:
+    if not len(source_to_replacement_dict):
+        return cg
+
+    source_to_replacement_dict = gamla.pipe(
+        source_to_replacement_dict,
+        gamla.keymap(graph.make_computation_node),
+        gamla.valmap(
+            gamla.unless(
+                gamla.is_instance(base_types.GraphType), graph.make_computation_node
+            )
+        ),
+    )
+    get_duplicate = gamla.pipe(
+        gamla.graph_traverse_many(
+            source_to_replacement_dict.keys(), graph.traverse_forward(cg.edges)
+        ),
+        gamla.remove(gamla.contains(source_to_replacement_dict.keys())),
+        _node_to_duplicated_node,
+    )
+
+    def node_replacement(node):
+        replacement = source_to_replacement_dict.get(node)
+        if replacement is not None:
+            return replacement
+        duplicate = get_duplicate(node)
+        return duplicate if duplicate is not None else node
+
+    replacement_graphs_edges = {}
+
+    def source_replacement(node):
+        replacement = node_replacement(node)
+        if isinstance(replacement, base_types.GraphType):
+            replacement_graphs_edges[id(replacement)] = replacement.edges
+            return replacement.sink
+        return replacement
+
+    new_edges = []
+    for edge in cg.edges:
+        destination = node_replacement(edge.destination)
+        if edge.args:
+            args = tuple(source_replacement(arg) for arg in edge.args)
+            if destination is edge.destination and all(
+                new is old for new, old in zip(args, edge.args)
+            ):
+                new_edges.append(edge)
+            else:
+                new_edges.append(
+                    dataclasses.replace(edge, args=args, destination=destination)
+                )
+        else:
+            source = source_replacement(edge.source)
+            if source is edge.source and destination is edge.destination:
+                new_edges.append(edge)
+            else:
+                new_edges.append(
+                    dataclasses.replace(edge, source=source, destination=destination)
+                )
+
+    return base_types.GraphType(
+        edges=base_types.merge_edges(
+            frozenset(new_edges), *replacement_graphs_edges.values()
+        ),
+        sink=node_replacement(cg.sink),
+    )

{computation_graph-67 → computation_graph-68}/computation_graph/graph.py

@@ -1,6 +1,6 @@
 import dataclasses
 import functools
-from typing import Callable, FrozenSet, Optional, Tuple
+from typing import Callable, Dict, FrozenSet, Optional, Tuple
 
 import gamla
 from gamla.optimized import sync as opt_gamla
@@ -231,6 +231,45 @@ def replace_node(
     )
 
 
+def replace_nodes(
+    node_map: Dict[base_types.ComputationNode, base_types.ComputationNode]
+) -> Callable[[base_types.GraphType], base_types.GraphType]:
+    """Replace many nodes at once (sources, destinations and sink) in a single pass."""
+
+    def replace_nodes(g: base_types.GraphType) -> base_types.GraphType:
+        if not node_map:
+            return g
+        new_edges = []
+        for edge in g.edges:
+            destination = node_map.get(edge.destination, edge.destination)
+            if edge.args:
+                args = tuple(node_map.get(arg, arg) for arg in edge.args)
+                if destination is edge.destination and all(
+                    new is old for new, old in zip(args, edge.args)
+                ):
+                    new_edges.append(edge)
+                else:
+                    new_edges.append(
+                        dataclasses.replace(edge, args=args, destination=destination)
+                    )
+            else:
+                assert edge.source is not None
+                source = node_map.get(edge.source, edge.source)
+                if source is edge.source and destination is edge.destination:
+                    new_edges.append(edge)
+                else:
+                    new_edges.append(
+                        dataclasses.replace(
+                            edge, source=source, destination=destination
+                        )
+                    )
+        return base_types.GraphType(
+            edges=frozenset(new_edges), sink=node_map.get(g.sink, g.sink)
+        )
+
+    return replace_nodes
+
+
 def transform_edges(
     query: Callable[[base_types.ComputationEdge], bool],
     edge_mapper: Callable[[base_types.ComputationEdge], base_types.ComputationEdge],

{computation_graph-67 → computation_graph-68}/computation_graph/graph_test.py

@@ -1000,6 +1000,33 @@ def test_replace_node():
     )
 
 
+@pytest.mark.parametrize(
+    "node_map",
+    [
+        pytest.param({}, id="empty map is a no-op"),
+        pytest.param({_node1: _node1_async}, id="single node"),
+        pytest.param({_node1: _node1_async, _node2: _next_int}, id="multiple nodes"),
+        pytest.param({(lambda x: x): _node1_async}, id="node not in graph is a no-op"),
+    ],
+)
+def test_replace_nodes(node_map: Dict):
+    a = graph.make_source()
+    g = graph.merge_graphs(
+        composers.compose_source_unary(_node1, a),
+        composers.compose_left_unary(_node1, _node2),
+        sink_node_or_graph=graph.make_computation_node(_node2),
+    )
+    node_map = gamla.pipe(
+        node_map,
+        gamla.keymap(graph.make_computation_node),
+        gamla.valmap(graph.make_computation_node),
+    )
+    expected = g
+    for original, replacement in node_map.items():
+        expected = graph.replace_node(original, replacement)(expected)
+    assert graph.replace_nodes(node_map)(g) == expected
+
+
 def test_ambig_edges_assertion_in_merge_graphs_active_only_when_env_var_is_active(
     monkeypatch,
 ):

{computation_graph-67 → computation_graph-68}/computation_graph.egg-info/PKG-INFO

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

{computation_graph-67 → computation_graph-68}/setup.py

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

computation_graph-67/computation_graph/composers/duplication.py

@@ -1,124 +0,0 @@
-import asyncio
-import functools
-import inspect
-from typing import Dict
-
-import gamla
-
-from computation_graph import base_types, graph
-from computation_graph.composers import debug
-
-
-def duplicate_function(func):
-    if asyncio.iscoroutinefunction(func):
-
-        @functools.wraps(func)
-        async def inner(*args, **kwargs):
-            return await func(*args, **kwargs)
-
-    else:
-
-        @functools.wraps(func)
-        def inner(*args, **kwargs):
-            return func(*args, **kwargs)
-
-    return debug.name_callable(inner, f"duplicate of {func.__name__}")
-
-
-def _duplicate_computation_edge(get_duplicated_node):
-    return gamla.compose_left(
-        gamla.dataclass_transform("source", get_duplicated_node),
-        gamla.dataclass_transform("destination", get_duplicated_node),
-        gamla.dataclass_transform(
-            "args", gamla.compose_left(gamla.map(get_duplicated_node), tuple)
-        ),
-    )
-
-
-_duplicate_node = gamla.compose_left(
-    base_types.node_implementation,
-    gamla.when(
-        gamla.compose_left(inspect.signature, gamla.attrgetter("parameters"), len),
-        duplicate_function,
-    ),
-    graph.make_computation_node,
-)
-
-duplicate_node = _duplicate_node
-
-_node_to_duplicated_node = gamla.compose_left(
-    gamla.remove(base_types.node_is_terminal),
-    gamla.map(gamla.pair_right(_duplicate_node)),
-    dict,
-    gamla.dict_to_getter_with_default(None),
-)
-
-
-def duplicate_graph(original: base_types.GraphType) -> base_types.GraphType:
-
-    _node_to_duplicated_node_index = gamla.pipe(
-        original.edges, graph.get_all_nodes, _node_to_duplicated_node
-    )
-    get_duplicated_node = gamla.when(
-        _node_to_duplicated_node_index, _node_to_duplicated_node_index
-    )
-    return base_types.GraphType(
-        edges=gamla.pipe(
-            original.edges,
-            gamla.map(
-                _duplicate_computation_edge(
-                    gamla.when(get_duplicated_node, get_duplicated_node)
-                )
-            ),
-            frozenset,
-        ),
-        sink=get_duplicated_node(original.sink),
-    )
-
-
-duplicate_function_or_graph = gamla.ternary(
-    gamla.is_instance(base_types.GraphType), 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:
-    if not len(source_to_replacement_dict):
-        return cg
-
-    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.edges)
-        ),
-        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))(
-            base_types.GraphType(edges=frozenset((e,)), sink=dest)
-        )
-        for source in base_types.edge_sources(e):
-            g = graph.replace_source(source, get_node_replacement(source), g)
-        return g
-
-    return base_types.GraphType(
-        edges=frozenset(edge for e in cg.edges for edge in update_edge(e).edges),
-        sink=get_node_replacement(cg.sink),
-    )