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

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

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

{computation_graph-66 → computation_graph-68}/computation_graph/base_types.py

@@ -63,9 +63,17 @@ class ComputationEdge:
         assert bool(self.args) != bool(
             self.source
         ), f"Edge must have a source or args, not both: {self}"
+        assert bool(self.args) or isinstance(
+            self.source, ComputationNode
+        ), f"source must be a ComputationNode {self.source}"
+        assert isinstance(
+            self.destination, ComputationNode
+        ), f"destination must be a ComputationNode {self.destination}"
+        assert all(
+            isinstance(x, ComputationNode) for x in self.args
+        ), f"all args must be ComputationNodes {self.args}"
         if (
             not self.args
-            # TODO(uri): doesn't support `functools.partial`, suggested to drop support for it entirely.
             and not isinstance(self.source.func, functools.partial)
             and not isinstance(self.destination.func, functools.partial)
         ):
@@ -76,7 +84,9 @@ class ComputationEdge:
 
     def __repr__(self):
         source_str = (
-            "".join(map(str, self.args)) if self.source is None else str(self.source)
+            f"[{''.join(map(str, self.args))}]"
+            if self.source is None
+            else str(self.source)
         )
         line = "...." if self.is_future else "----"
         return source_str + line + self.key + line + ">" + str(self.destination)
@@ -105,7 +115,16 @@ class ComputationNode:
         return self.computed_hash
 
     def __repr__(self):
-        return self.name
+        return f"<CompuationNode {self.name.replace('<', '').replace('>', '')}({','.join(self.signature.kwargs)}) >"
+
+
+@dataclasses.dataclass(frozen=True)
+class GraphType:
+    edges: FrozenSet[ComputationEdge]
+    sink: ComputationNode
+
+    def __post_init__(self):
+        assert isinstance(self.edges, frozenset), "edges must be a frozenset"
 
 
 node_implementation = gamla.attrgetter("func")
@@ -125,13 +144,15 @@ def edge_sources(edge: ComputationEdge) -> Tuple[ComputationNode, ...]:
 
 is_computation_graph = gamla.alljuxt(
     # Note that this cannot be set to `GraphType` (due to `is_instance` limitation).
-    gamla.is_instance((tuple, set, frozenset)),
-    gamla.len_greater(0),
-    gamla.allmap(gamla.is_instance(ComputationEdge)),
+    gamla.is_instance(GraphType),
+    gamla.compose_left(gamla.attrgetter("edges"), gamla.len_greater(0)),
+    gamla.compose_left(
+        gamla.attrgetter("edges"), gamla.allmap(gamla.is_instance(ComputationEdge))
+    ),
 )
 
 
-ambiguity_groups = gamla.compose(
+_ambiguity_groups_on_edges = gamla.compose(
     frozenset,
     gamla.filter(gamla.len_greater(1)),
     dict.values,
@@ -139,9 +160,9 @@ ambiguity_groups = gamla.compose(
     gamla.groupby(gamla.juxt(edge_destination, edge_key, edge_priority)),
 )
 
-assert_no_unwanted_ambiguity = gamla.side_effect(
+assert_no_unwanted_ambiguity_on_edges = gamla.side_effect(
     gamla.compose_left(
-        ambiguity_groups,
+        _ambiguity_groups_on_edges,
         gamla.assert_that_with_message(
             gamla.wrap_str(
                 "There are multiple edges with the same destination, key and priority in the computation graph!: {}"
@@ -153,23 +174,24 @@ assert_no_unwanted_ambiguity = gamla.side_effect(
 )
 
 
-@gamla.side_effect
-def _assert_no_unwanted_ambiguity_when_debug_set(graph):
-    if os.getenv(COMPUTATION_GRAPH_DEBUG_ENV_KEY) is not None:
-        assert_no_unwanted_ambiguity(graph)
+def is_debug_mode() -> bool:
+    return os.getenv(COMPUTATION_GRAPH_DEBUG_ENV_KEY) is not None
 
 
-def merge_graphs(*graphs):
-    s = frozenset({})
-    s = s.union(*graphs)
+@gamla.side_effect
+def assert_no_unwanted_ambiguity_when_debug_set(graph: GraphType):
+    if is_debug_mode():
+        assert_no_unwanted_ambiguity_on_edges(graph.edges)
+
 
-    return _assert_no_unwanted_ambiguity_when_debug_set(s)
+# TODO(nitzo): Is this used ? Check if we can re-enable the assertion when in debug / move to merge.
+def merge_edges(*edges: FrozenSet[ComputationEdge]) -> FrozenSet[ComputationEdge]:
+    result: FrozenSet[ComputationEdge] = frozenset()
+    return result.union(*edges)
 
 
-# We use a tuple to generate a unique id for each node based on the order of edges.
-GraphType = FrozenSet[ComputationEdge]
 GraphOrCallable = Union[Callable, GraphType]
 CallableOrNode = Union[Callable, ComputationNode]
 CallableOrNodeOrGraph = Union[CallableOrNode, GraphType]
 NodeOrGraph = Union[ComputationNode, GraphType]
-EMPTY_GRAPH: GraphType = frozenset()
+EMPTY_GRAPH: GraphType = GraphType(edges=frozenset(), sink=None)  # type: ignore[arg-type]

{computation_graph-66 → computation_graph-68}/computation_graph/composers/__init__.py

@@ -1,18 +1,20 @@
 from __future__ import annotations
 
-from typing import Any, Callable, Dict, Iterable, Optional, Sequence
+from typing import Any, Callable, Dict, Iterable, Optional, Sequence, Tuple
 
 import gamla
 
 from computation_graph import base_types, graph, signature
+from computation_graph.base_types import GraphType
+from computation_graph.graph import make_computation_node
 
 
 class _ComputationError:
     pass
 
 
-_callable_or_graph_type_to_node_or_graph_type = gamla.unless(
-    gamla.is_instance((tuple, set, frozenset)), graph.make_computation_node
+_to_computation_node_if_callable = gamla.unless(
+    gamla.is_instance(base_types.GraphType), graph.make_computation_node
 )
 
 
@@ -47,37 +49,52 @@ def make_and(
         return args
 
     merge_node = graph.make_computation_node(args_to_tuple)
+    # `merge_fn` may be a callable/node or a whole graph; `make_compose` handles
+    # both, connecting `merge_node`'s output into `merge_fn` under key "args".
+    merge_fn_graph = make_compose(merge_fn, merge_node, key="args")
 
     return gamla.sync.pipe(
         funcs,
-        gamla.sync.map(_callable_or_graph_type_to_node_or_graph_type),
+        gamla.sync.map(_to_computation_node_if_callable),
         tuple,
         gamla.sync.juxtcat(
-            gamla.sync.map(_get_edges_from_node_or_graph),
+            gamla.sync.filter(gamla.is_instance(base_types.GraphType)),
             gamla.sync.compose_left(
-                gamla.sync.map(_infer_sink),
+                gamla.sync.map(
+                    lambda node_or_graph: (
+                        node_or_graph.sink
+                        if isinstance(node_or_graph, base_types.GraphType)
+                        else node_or_graph
+                    )
+                ),
                 tuple,
-                lambda nodes: (
-                    (
-                        base_types.ComputationEdge(
-                            is_future=False,
-                            priority=0,
-                            source=None,
-                            args=nodes,
-                            destination=merge_node,
-                            key="*args",
-                        ),
+                lambda nodes: base_types.GraphType(
+                    frozenset(
+                        [
+                            base_types.ComputationEdge(
+                                is_future=False,
+                                priority=0,
+                                source=None,
+                                args=nodes,
+                                destination=merge_node,
+                                key="*args",
+                            )
+                        ]
                     ),
-                    make_compose(merge_fn, merge_node, key="args"),
+                    merge_node,
                 ),
+                gamla.wrap_tuple,
             ),
         ),
-        gamla.sync.star(base_types.merge_graphs),
+        tuple,
+        lambda graphs: graph.merge_graphs(
+            *graphs, merge_fn_graph, sink_node_or_graph=merge_fn_graph
+        ),
     )
 
 
 def make_or(
-    funcs: Sequence[base_types.CallableOrNodeOrGraph],
+    nodes_or_graphs: Sequence[base_types.CallableOrNodeOrGraph],
     merge_fn: base_types.CallableOrNodeOrGraph,
 ) -> base_types.GraphType:
     """Aggregate funcs' output into `merge_fn`.
@@ -91,63 +108,44 @@ def make_or(
         )
 
     filter_node = graph.make_computation_node(filter_computation_errors)
+    # `merge_fn` may be a callable/node or a whole graph; `make_compose` handles
+    # both, connecting `filter_node`'s output into `merge_fn` under key "args".
+    merge_fn_graph = make_compose(merge_fn, filter_node, key="args")
 
     return gamla.sync.pipe(
-        funcs,
+        nodes_or_graphs,
         gamla.sync.map(make_optional(default_value=_ComputationError())),
         tuple,
-        gamla.sync.pair_right(
+        gamla.juxtcat(
+            gamla.identity,
             gamla.sync.compose_left(
-                gamla.sync.map(_infer_sink),
+                gamla.sync.map(gamla.attrgetter("sink")),
                 tuple,
-                lambda sinks: (
-                    (
-                        base_types.ComputationEdge(
-                            is_future=False,
-                            priority=0,
-                            source=None,
-                            args=sinks,
-                            destination=filter_node,
-                            key="*args",
-                        ),
+                lambda sinks: GraphType(
+                    frozenset(
+                        [
+                            base_types.ComputationEdge(
+                                is_future=False,
+                                priority=0,
+                                source=None,
+                                args=sinks,
+                                destination=filter_node,
+                                key="*args",
+                            )
+                        ]
                     ),
-                    make_compose(merge_fn, filter_node, key="args"),
+                    filter_node,
                 ),
-            )
+                gamla.wrap_tuple,
+            ),
+        ),
+        tuple,
+        lambda graphs: graph.merge_graphs(
+            *graphs, merge_fn_graph, sink_node_or_graph=merge_fn_graph
         ),
-        gamla.concat,
-        gamla.sync.star(base_types.merge_graphs),
     )
 
 
-_destinations = gamla.compose(set, gamla.map(base_types.edge_destination))
-
-
-def _infer_sink(graph_or_node: base_types.NodeOrGraph) -> base_types.ComputationNode:
-    if isinstance(graph_or_node, base_types.ComputationNode):
-        return graph_or_node
-    graph_without_future_edges = graph.remove_future_edges(graph_or_node)
-    if graph_without_future_edges:
-        try:
-            return graph.sink_excluding_terminals(graph_without_future_edges)
-        except AssertionError:
-            # If we reached here we can try again without sources of future edges.
-            sources_of_future_edges = gamla.sync.pipe(
-                graph_or_node,
-                gamla.sync.filter(base_types.edge_is_future),
-                gamla.sync.map(base_types.edge_source),
-                frozenset,
-            )
-            result = (
-                graph.get_leaves(graph_without_future_edges) - sources_of_future_edges
-            )
-            assert len(result) == 1
-            return gamla.head(result)
-
-    assert len(_destinations(graph_or_node)) == 1, graph_or_node
-    return base_types.edge_destination(gamla.head(graph_or_node))
-
-
 def make_first(*graphs: base_types.CallableOrNodeOrGraph) -> base_types.GraphType:
     """Returns a graph that when run, returns the first value that doesn't raise an exception.
     >>> def raise_some_exception():
@@ -156,22 +154,34 @@ def make_first(*graphs: base_types.CallableOrNodeOrGraph) -> base_types.GraphTyp
     (raise_exception----constituent_of_first---->first_sink, just----constituent_of_first---->first_sink)
     Will return 1.
     """
-    graph_or_nodes = tuple(map(_callable_or_graph_type_to_node_or_graph_type, graphs))
+    graph_or_nodes = tuple(map(_to_computation_node_if_callable, graphs))
 
     def first_sink(constituent_of_first):
         return constituent_of_first
 
-    return base_types.merge_graphs(
-        *map(_get_edges_from_node_or_graph, graph_or_nodes),
+    sink = graph.make_computation_node(first_sink)
+
+    return graph.merge_graphs(
+        *gamla.pipe(
+            graph_or_nodes,
+            gamla.map(_get_edges_from_node_or_graph),
+            gamla.filter(gamla.is_instance(GraphType)),
+        ),
         gamla.pipe(
             graph_or_nodes,
-            gamla.map(_infer_sink),
+            gamla.map(
+                lambda graph_or_node: (
+                    graph_or_node.sink
+                    if isinstance(graph_or_node, GraphType)
+                    else graph_or_node
+                )
+            ),
             enumerate,
             gamla.map(
                 gamla.star(
                     lambda i, g: base_types.ComputationEdge(
                         source=g,
-                        destination=graph.make_computation_node(first_sink),
+                        destination=sink,
                         key="constituent_of_first",
                         args=(),
                         priority=i,
@@ -179,8 +189,10 @@ def make_first(*graphs: base_types.CallableOrNodeOrGraph) -> base_types.GraphTyp
                     )
                 )
             ),
-            tuple,
+            frozenset,
+            lambda edges: GraphType(edges, sink),
         ),
+        sink_node_or_graph=sink,
     )
 
 
@@ -188,6 +200,39 @@ def last(*args) -> base_types.GraphType:
     return make_first(*reversed(args))
 
 
+def _determine_sink(
+    source: base_types.NodeOrGraph, destination: base_types.NodeOrGraph, is_future: bool
+) -> base_types.ComputationNode:
+    if is_future:
+        if isinstance(source, base_types.GraphType):
+            return source.sink
+        return (
+            destination.sink
+            if isinstance(destination, base_types.GraphType)
+            else destination
+        )
+
+    if isinstance(destination, base_types.GraphType):
+        return destination.sink
+
+    if destination.is_terminal:
+        return source.sink if isinstance(source, base_types.GraphType) else source
+    return destination
+
+
+def _determine_composition_sink(
+    graphs: Tuple[base_types.NodeOrGraph, ...]
+) -> base_types.NodeOrGraph:
+    for graph_or_node in reversed(graphs):
+        if isinstance(graph_or_node, base_types.GraphType):
+            return graph_or_node.sink
+
+        if graph_or_node.is_terminal:
+            continue
+
+    raise AssertionError("Could not compose, failed to. find a sink")
+
+
 @gamla.curry
 def _try_connect(
     source: base_types.ComputationNode,
@@ -220,20 +265,31 @@ def _infer_composition_edges(
     source: base_types.NodeOrGraph,
     destination: base_types.NodeOrGraph,
 ) -> base_types.GraphType:
-    try_connect = _try_connect(_infer_sink(source), key, priority, is_future)
+
+    try_connect = _try_connect(
+        source.sink if isinstance(source, GraphType) else source,
+        key,
+        priority,
+        is_future,
+    )
 
     if isinstance(destination, base_types.ComputationNode):
-        return base_types.merge_graphs(
-            (try_connect(destination, destination.signature),),
+        sink = _determine_sink(source, destination, is_future)
+        return graph.merge_graphs(
+            base_types.GraphType(
+                edges=frozenset([try_connect(destination, destination.signature)]),
+                sink=sink,
+            ),
             _get_edges_from_node_or_graph(source),
+            sink_node_or_graph=sink,
         )
-
     unbound_signature = graph.unbound_signature(
-        graph.get_incoming_edges_for_node(destination)
+        graph.get_incoming_edges_for_node(destination.edges)
     )
-    return base_types.merge_graphs(
+    return graph.merge_graphs(
         gamla.sync.pipe(
             destination,
+            gamla.attrgetter("edges"),
             gamla.sync.mapcat(graph.get_edge_nodes),
             gamla.unique,
             gamla.sync.filter(
@@ -246,24 +302,31 @@ def _infer_composition_edges(
             # 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)
+                or node
+                not in graph.get_all_nodes(
+                    source.edges if isinstance(source, GraphType) else source
+                )
             ),
             gamla.sync.map(
-                lambda destination: try_connect(
-                    destination=destination,
-                    unbound_destination_signature=unbound_signature(destination),
+                lambda destination_node: try_connect(
+                    destination=destination_node,
+                    unbound_destination_signature=unbound_signature(destination_node),
                 )
             ),
-            tuple,
+            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]
         ),
-        destination,
         _get_edges_from_node_or_graph(source),
+        destination,
+        sink_node_or_graph=_determine_sink(source, destination, is_future),
     )
 
 
@@ -279,10 +342,13 @@ def _make_compose_inner(
     return gamla.sync.pipe(
         funcs,
         reversed,
-        gamla.sync.map(_callable_or_graph_type_to_node_or_graph_type),
+        gamla.sync.map(_to_computation_node_if_callable),
         gamla.sliding_window(2),
         gamla.sync.map(gamla.star(_infer_composition_edges(priority, key, is_future))),
-        gamla.sync.star(base_types.merge_graphs),
+        tuple,
+        lambda graphs: graph.merge_graphs(
+            *graphs, sink_node_or_graph=_determine_composition_sink(graphs)
+        ),
     )
 
 
@@ -309,11 +375,24 @@ def make_compose_future(
     def when_memory_unavailable():
         return default
 
-    return base_types.merge_graphs(
+    if not isinstance(destination, base_types.GraphType):
+        destination = make_computation_node(destination)
+
+    if not isinstance(source, base_types.GraphType):
+        source = make_computation_node(source)
+
+    if isinstance(source, base_types.GraphType):
+        sink_node_or_graph = source.sink
+    elif isinstance(destination, base_types.GraphType):
+        sink_node_or_graph = destination.sink
+    else:
+        sink_node_or_graph = destination
+    return graph.merge_graphs(
         _make_compose_inner(destination, source, key=key, is_future=True, priority=0),
         _make_compose_inner(
             destination, when_memory_unavailable, key=key, is_future=False, priority=1
         ),
+        sink_node_or_graph=sink_node_or_graph,
     )
 
 
@@ -384,11 +463,22 @@ def compose_dict(
     >>> compose_dict(gamla.between, {"low": gamla.just(0), "high": gamla.just(10)})
     (just----low---->between, just----high---->between)
     """
+
+    destination = f if base_types.is_computation_graph(f) else make_computation_node(f)  # type: ignore
+    sink = (
+        destination.sink  # type: ignore
+        if base_types.is_computation_graph(destination)
+        else destination
+    )
+
     return gamla.pipe(
         d,
         dict.items,
-        gamla.sync.map(gamla.star(lambda key, fn: make_compose(f, fn, key=key))),
-        gamla.sync.star(base_types.merge_graphs),
+        gamla.sync.map(
+            gamla.star(lambda key, fn: make_compose(destination, fn, key=key))
+        ),
+        tuple,
+        lambda graphs: graph.merge_graphs(*graphs, sink_node_or_graph=sink),  # type: ignore
     ) or compose_left_unary(f, lambda x: x)
 
 

{computation_graph-66 → computation_graph-68}/computation_graph/composers/composers_test.py

@@ -1,5 +1,5 @@
 import computation_graph.graph
-from computation_graph import base_types, composers, graph_runners
+from computation_graph import composers, graph, graph_runners
 
 
 def test_unary_composition_with_graph_destination():
@@ -30,14 +30,36 @@ def test_infer_sink_edge_case_all_future_edges_with_single_destination():
     def c_b(c, b):
         return b
 
-    two_future_edges_single_dest = base_types.merge_graphs(
+    two_future_edges_single_dest = graph.merge_graphs(
         composers.compose_left_source(s, "c", c_b),
         composers.compose_left_source(s, "b", c_b),
+        sink_node_or_graph=graph.make_computation_node(c_b),
     )
 
     composers.compose_left_unary(two_future_edges_single_dest, lambda x: x)
 
 
+def test_compose_dict():
+    def add(a, b):
+        return a + b
+
+    result = graph_runners.nullary_infer_sink(
+        composers.compose_dict(add, {"a": lambda: 3, "b": lambda: 4})
+    )
+    assert result == 7
+
+
+def test_compose_dict_with_graph_destination():
+    def add(a, b):
+        return a + b
+
+    destination_graph = composers.compose_left(lambda: 3, add, key="a")
+    result = graph_runners.nullary_infer_sink(
+        composers.compose_dict(destination_graph, {"b": lambda: 4})
+    )
+    assert result == 7
+
+
 def test_ambiguity_does_not_blow_up():
     counter = 0