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

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

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

{computation_graph-67 → computation_graph-69}/computation_graph/composers/__init__.py

@@ -286,44 +286,36 @@ def _infer_composition_edges(
     unbound_signature = graph.unbound_signature(
         graph.get_incoming_edges_for_node(destination.edges)
     )
+    source_nodes = (
+        graph.get_all_nodes(source.edges)
+        if isinstance(source, GraphType)
+        else frozenset()
+    )
+    skip_cycle_check = isinstance(source, base_types.ComputationNode)
+    new_edges = set()
+    # The connectable-node set is order-independent (result is a frozenset), so we
+    # can iterate the cached node set instead of rescanning the edges.
+    for node in graph.get_all_nodes(destination.edges):
+        node_signature = unbound_signature(node)
+        if not (
+            key in node_signature.kwargs
+            or (key is None and signature.is_unary(node_signature))
+        ):
+            continue
+        # Do not add edges to nodes from source that are already present in destination (cycle).
+        if not skip_cycle_check and node in source_nodes:
+            continue
+        new_edges.add(
+            try_connect(destination=node, unbound_destination_signature=node_signature)
+        )
+    if not new_edges:
+        raise AssertionError(
+            f"Cannot compose, destination signature does not contain key '{key}'"
+        )
     return graph.merge_graphs(
-        gamla.sync.pipe(
-            destination,
-            gamla.attrgetter("edges"),
-            gamla.sync.mapcat(graph.get_edge_nodes),
-            gamla.unique,
-            gamla.sync.filter(
-                gamla.sync.compose_left(
-                    unbound_signature,
-                    lambda sig: key in sig.kwargs
-                    or (key is None and signature.is_unary(sig)),
-                )
-            ),
-            # Do not add edges to nodes from source that are already present in destination (cycle).
-            gamla.sync.filter(
-                lambda node: isinstance(source, base_types.ComputationNode)
-                or node
-                not in graph.get_all_nodes(
-                    source.edges if isinstance(source, GraphType) else source
-                )
-            ),
-            gamla.sync.map(
-                lambda destination_node: try_connect(
-                    destination=destination_node,
-                    unbound_destination_signature=unbound_signature(destination_node),
-                )
-            ),
-            frozenset,
-            gamla.sync.check(
-                gamla.identity,
-                AssertionError(
-                    f"Cannot compose, destination signature does not contain key '{key}'"
-                ),
-            ),
-            # Sink is a placeholder here; `merge_graphs` below sets the real sink
-            # via `sink_node_or_graph`.
-            lambda edges: GraphType(edges, None),  # type: ignore[arg-type]
-        ),
+        # Sink is a placeholder here; `merge_graphs` below sets the real sink
+        # via `sink_node_or_graph`.
+        GraphType(frozenset(new_edges), None),  # type: ignore[arg-type]
         _get_edges_from_node_or_graph(source),
         destination,
         sink_node_or_graph=_determine_sink(source, destination, is_future),

computation_graph-69/computation_graph/composers/duplication.py

@@ -0,0 +1,154 @@
+import dataclasses
+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 inspect.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-69}/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, Iterable, Optional, Tuple
 
 import gamla
 from gamla.optimized import sync as opt_gamla
@@ -64,11 +64,17 @@ def make_computation_node(
     )
 
 
-get_incoming_edges_for_node = opt_gamla.compose_left(
-    opt_gamla.groupby(base_types.edge_destination),
-    opt_gamla.valmap(frozenset),
-    gamla.dict_to_getter_with_default(frozenset()),
-)
+_EMPTY_EDGE_SET: FrozenSet[base_types.ComputationEdge] = frozenset()
+
+
+def get_incoming_edges_for_node(
+    edges: Iterable[base_types.ComputationEdge],
+) -> Callable[[base_types.ComputationNode], FrozenSet[base_types.ComputationEdge]]:
+    grouped: Dict[base_types.ComputationNode, list] = {}
+    for edge in edges:
+        grouped.setdefault(edge.destination, []).append(edge)
+    node_to_edges = {node: frozenset(group) for node, group in grouped.items()}
+    return lambda node: node_to_edges.get(node, _EMPTY_EDGE_SET)
 
 
 get_terminals = opt_gamla.compose_left(
@@ -95,30 +101,24 @@ def make_terminal(name: str, func: Callable) -> base_types.ComputationNode:
     )
 
 
-_keep_not_in_bound_kwargs = opt_gamla.compose_left(
-    opt_gamla.map(base_types.edge_key),
-    gamla.filter(gamla.identity),
-    frozenset,
-    gamla.contains,
-    gamla.remove,
-)
-
-
 @gamla.curry
 def unbound_signature(
     node_to_incoming_edges, node: base_types.ComputationNode
 ) -> base_types.NodeSignature:
     """Computes the new signature of unbound variables after considering internal edges."""
     incoming_edges = node_to_incoming_edges(node)
-    keep_not_in_bound_kwargs = _keep_not_in_bound_kwargs(incoming_edges)
+    bound_keys = {edge.key for edge in incoming_edges if edge.key}
+    node_signature = node.signature
     return base_types.NodeSignature(
-        is_kwargs=node.signature.is_kwargs
-        and "**kwargs" not in tuple(opt_gamla.map(base_types.edge_key)(incoming_edges))
-        and signature.is_unary(node.signature),
-        is_args=node.signature.is_args
+        is_kwargs=node_signature.is_kwargs
+        and all(edge.key != "**kwargs" for edge in incoming_edges)
+        and signature.is_unary(node_signature),
+        is_args=node_signature.is_args
         and not any(edge.args for edge in incoming_edges),
-        kwargs=tuple(keep_not_in_bound_kwargs(node.signature.kwargs)),
-        optional_kwargs=tuple(keep_not_in_bound_kwargs(node.signature.optional_kwargs)),
+        kwargs=tuple(k for k in node_signature.kwargs if k not in bound_keys),
+        optional_kwargs=tuple(
+            k for k in node_signature.optional_kwargs if k not in bound_keys
+        ),
     )
 
 
@@ -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],
@@ -322,28 +361,20 @@ def merge_graphs(
         isinstance(x, base_types.GraphType) for x in graphs
     ), f"{graphs} not a graph"
     assert (
-        not base_types.is_computation_graph(sink_node_or_graph)
+        not isinstance(sink_node_or_graph, base_types.GraphType)
+        or not sink_node_or_graph.edges
         or sink_node_or_graph in graphs
     ), "sink graph must be one of the graphs being merged"
 
-    new_g = gamla.pipe(
-        graphs,
-        gamla.remove(gamla.equals(base_types.EMPTY_GRAPH)),
-        tuple,
-        gamla.pair_right(
-            gamla.just(
-                sink_node_or_graph.sink
-                if isinstance(sink_node_or_graph, base_types.GraphType)
-                else sink_node_or_graph
-            )
+    new_g = base_types.GraphType(
+        base_types.merge_edges(
+            *(g.edges for g in graphs if g != base_types.EMPTY_GRAPH)
         ),
-        gamla.packstack(
-            gamla.compose_left(
-                gamla.map(gamla.attrgetter("edges")), gamla.star(base_types.merge_edges)
-            ),
-            gamla.identity,
+        (
+            sink_node_or_graph.sink
+            if isinstance(sink_node_or_graph, base_types.GraphType)
+            else sink_node_or_graph
         ),
-        gamla.star(base_types.GraphType),
     )
 
     if base_types.is_debug_mode():

{computation_graph-67 → computation_graph-69}/computation_graph/graph_runners.py

@@ -1,4 +1,4 @@
-import asyncio
+import inspect
 from typing import Callable, Union
 
 import gamla
@@ -77,7 +77,7 @@ def unary_with_state_and_expectations(
 def variadic_with_state_and_expectations(g, sink):
     f = run.to_callable_strict(g)
 
-    if asyncio.iscoroutinefunction(f):
+    if inspect.iscoroutinefunction(f):
 
         async def inner(turns):
             prev = {}

{computation_graph-67 → computation_graph-69}/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-69}/computation_graph/run.py

@@ -109,7 +109,7 @@ _is_graph_async = opt_gamla.compose_left(
     opt_gamla.mapcat(lambda edge: (edge.source, *edge.args)),
     opt_gamla.remove(gamla.equals(None)),
     opt_gamla.map(base_types.node_implementation),
-    gamla.anymap(asyncio.iscoroutinefunction),
+    gamla.anymap(inspect.iscoroutinefunction),
 )
 
 
@@ -452,7 +452,7 @@ def _make_get_node_executor(
         return result
 
     all_nodes = graph.get_all_nodes(edges)
-    async_nodes = {n for n in all_nodes if asyncio.iscoroutinefunction(n.func)}
+    async_nodes = {n for n in all_nodes if inspect.iscoroutinefunction(n.func)}
     sync = all_nodes - async_nodes
     tf = graph.traverse_forward(edges)
     downstream_from_async = set(gamla.graph_traverse_many(async_nodes, tf))
@@ -572,7 +572,9 @@ to_callable = to_callable_with_side_effect(gamla.just(gamla.just(None)))
 
 
 def _node_is_properly_composed(
-    node_to_incoming_edges: base_types.GraphType,
+    node_to_incoming_edges: Callable[
+        [base_types.ComputationNode], FrozenSet[base_types.ComputationEdge]
+    ]
 ) -> Callable[[base_types.ComputationNode], bool]:
     return gamla.compose_left(
         graph.unbound_signature(node_to_incoming_edges),
@@ -581,7 +583,7 @@ def _node_is_properly_composed(
     )
 
 
-def _assert_composition_is_valid(g: base_types.GraphType):
+def _assert_composition_is_valid(g: typing.Iterable[base_types.ComputationEdge]):
     return opt_gamla.pipe(
         g,
         graph.get_all_nodes,

{computation_graph-67 → computation_graph-69}/computation_graph/signature.py

@@ -1,5 +1,6 @@
 import functools
 import inspect
+import types
 from types import MappingProxyType
 from typing import Callable, FrozenSet
 
@@ -39,7 +40,7 @@ _func_parameters = gamla.compose_left(
 )
 
 
-def from_callable(func: Callable) -> base_types.NodeSignature:
+def _from_callable_via_inspect(func: Callable) -> base_types.NodeSignature:
     function_parameters = _func_parameters(func)
     return base_types.NodeSignature(
         is_args=gamla.anymap(parameter_is_star)(function_parameters),
@@ -60,6 +61,40 @@ def from_callable(func: Callable) -> base_types.NodeSignature:
     )
 
 
+def _as_plain_function(func: Callable):
+    """The `inspect.signature`-equivalent target when it is a plain function whose
+    signature is fully described by its code object, else None (wrappers with
+    `__signature__`, partials, methods, builtins and classes need `inspect`)."""
+    unwrapped = inspect.unwrap(func, stop=lambda f: hasattr(f, "__signature__"))
+    if not isinstance(unwrapped, types.FunctionType) or hasattr(
+        unwrapped, "__signature__"
+    ):
+        return None
+    return unwrapped
+
+
+def from_callable(func: Callable) -> base_types.NodeSignature:
+    plain_function = _as_plain_function(func)
+    if plain_function is None:
+        return _from_callable_via_inspect(func)
+    code = plain_function.__code__
+    positional = code.co_varnames[: code.co_argcount]
+    keyword_only = code.co_varnames[
+        code.co_argcount : code.co_argcount + code.co_kwonlyargcount
+    ]
+    defaults = plain_function.__defaults__ or ()
+    keyword_defaults = plain_function.__kwdefaults__ or {}
+    return base_types.NodeSignature(
+        is_args=bool(code.co_flags & inspect.CO_VARARGS),
+        is_kwargs=bool(code.co_flags & inspect.CO_VARKEYWORDS),
+        kwargs=positional + keyword_only,
+        optional_kwargs=(
+            positional[len(positional) - len(defaults) :]
+            + tuple(name for name in keyword_only if name in keyword_defaults)
+        ),
+    )
+
+
 def parameters(signature: base_types.NodeSignature) -> FrozenSet[str]:
     return frozenset(
         {

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

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

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

@@ -7,7 +7,7 @@ with open("README.md", "r") as fh:
 setuptools.setup(
     name="computation-graph",
     python_requires=">=3",
-    version="67",
+    version="69",
     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),
-    )