Issue #391 avoid multiple "result" modes with MultiResult flattening

Make GraphFlattener multi-input-sensitive, and be more careful about accidental mutation in this context

Author: soxofaan

openeo/internal/graph_building.py

@@ -10,6 +10,7 @@
 
 import abc
 import collections
+import copy
 import json
 import sys
 from contextlib import nullcontext
@@ -322,20 +323,29 @@ def generate(self, process_id: str):
 
 class GraphFlattener(ProcessGraphVisitor):
 
-    def __init__(self, node_id_generator: FlatGraphNodeIdGenerator = None):
+    def __init__(self, node_id_generator: FlatGraphNodeIdGenerator = None, multi_input_mode: bool = False):
         super().__init__()
         self._node_id_generator = node_id_generator or FlatGraphNodeIdGenerator()
         self._last_node_id = None
         self._flattened: Dict[str, dict] = {}
         self._argument_stack = []
         self._node_cache = {}
+        self._multi_input_mode = multi_input_mode
 
     def flatten(self, node: PGNode) -> Dict[str, dict]:
         """Consume given nested process graph and return flat dict representation"""
+        if self._flattened and not self._multi_input_mode:
+            raise RuntimeError("Flattening multiple graphs, but not in multi-input mode")
         self.accept_node(node)
         assert len(self._argument_stack) == 0
-        self._flattened[self._last_node_id]["result"] = True
-        return self._flattened
+        return self.flattened(set_result_flag=not self._multi_input_mode)
+
+    def flattened(self, set_result_flag: bool = True) -> Dict[str, dict]:
+        flat_graph = copy.deepcopy(self._flattened)
+        if set_result_flag:
+            # TODO #583 an "end" node is not necessarily a "result" node
+            flat_graph[self._last_node_id]["result"] = True
+        return flat_graph
 
     def accept_node(self, node: PGNode):
         # Process reused nodes only first time and remember node id.
@@ -450,14 +460,13 @@ def __init__(self, leaves: List[FlatGraphableMixin]):
         self._leaves = leaves
 
     def flat_graph(self) -> Dict[str, dict]:
-        result = {}
-        flattener = GraphFlattener()
+        flattener = GraphFlattener(multi_input_mode=True)
         for leaf in self._leaves:
             if isinstance(leaf, PGNode):
-                result = flattener.flatten(leaf)
+                flattener.flatten(leaf)
             elif isinstance(leaf, _FromNodeMixin):
-                result = flattener.flatten(leaf.from_node())
+                flattener.flatten(leaf.from_node())
             else:
-                raise ValueError(leaf)
+                raise ValueError(f"Unsupported type {type(leaf)}")
 
-        return result
+        return flattener.flattened(set_result_flag=True)

tests/internal/test_graphbuilding.py

@@ -8,6 +8,7 @@
 from openeo.api.process import Parameter
 from openeo.internal.graph_building import (
     FlatGraphNodeIdGenerator,
+    GraphFlattener,
     MultiResult,
     PGNode,
     PGNodeGraphUnflattener,
@@ -146,6 +147,91 @@ def test_flat_graph_key_generate():
     assert g.generate("foo") == "foo3"
 
 
+class TestGraphFlattener:
+    def test_simple(self):
+        node = PGNode("foo", bar="meh")
+        flattener = GraphFlattener()
+        assert flattener.flatten(node) == {"foo1": {"process_id": "foo", "arguments": {"bar": "meh"}, "result": True}}
+
+    def test_chain(self):
+        a = PGNode("a", bar="meh")
+        b = PGNode("b", a=a)
+        c = PGNode("c", a=a, b=b)
+        flattener = GraphFlattener()
+        assert flattener.flatten(c) == {
+            "a1": {"process_id": "a", "arguments": {"bar": "meh"}},
+            "b1": {"process_id": "b", "arguments": {"a": {"from_node": "a1"}}},
+            "c1": {
+                "process_id": "c",
+                "arguments": {"a": {"from_node": "a1"}, "b": {"from_node": "b1"}},
+                "result": True,
+            },
+        }
+
+    def test_no_multi_input_mode(self):
+        a = PGNode("a")
+        b = PGNode("b", a=a)
+        flattener = GraphFlattener()
+        flat_graph = flattener.flatten(a)
+        assert flat_graph == {"a1": {"process_id": "a", "arguments": {}, "result": True}}
+        with pytest.raises(RuntimeError, match="not in multi-input mode"):
+            flattener.flatten(b)
+        assert flat_graph == {"a1": {"process_id": "a", "arguments": {}, "result": True}}
+
+    def test_multi_input_mode(self):
+        a = PGNode("a")
+        b = PGNode("b", a=a)
+        c = PGNode("c", a=a)
+        flattener = GraphFlattener(multi_input_mode=True)
+        # Flatten b
+        assert flattener.flatten(b) == {
+            "a1": {"process_id": "a", "arguments": {}},
+            "b1": {"process_id": "b", "arguments": {"a": {"from_node": "a1"}}},
+        }
+        assert flattener.flattened() == {
+            "a1": {"process_id": "a", "arguments": {}},
+            "b1": {"process_id": "b", "arguments": {"a": {"from_node": "a1"}}, "result": True},
+        }
+        # Flatten c
+        assert flattener.flatten(c) == {
+            "a1": {"process_id": "a", "arguments": {}},
+            "b1": {"process_id": "b", "arguments": {"a": {"from_node": "a1"}}},
+            "c1": {"process_id": "c", "arguments": {"a": {"from_node": "a1"}}},
+        }
+        assert flattener.flattened() == {
+            "a1": {"process_id": "a", "arguments": {}},
+            "b1": {"process_id": "b", "arguments": {"a": {"from_node": "a1"}}},
+            "c1": {"process_id": "c", "arguments": {"a": {"from_node": "a1"}}, "result": True},
+        }
+
+    def test_multi_input_mode_mutation(self):
+        """Verify that previously produced flat graphs are not silently mutated"""
+        a = PGNode("a")
+        b = PGNode("b", a=a)
+        flattener = GraphFlattener(multi_input_mode=True)
+        a_flat = flattener.flatten(a)
+        assert a_flat == {
+            "a1": {"process_id": "a", "arguments": {}},
+        }
+        b_flat = flattener.flatten(b)
+        assert b_flat == {
+            "a1": {"process_id": "a", "arguments": {}},
+            "b1": {"process_id": "b", "arguments": {"a": {"from_node": "a1"}}},
+        }
+        assert flattener.flattened() == {
+            "a1": {"process_id": "a", "arguments": {}},
+            "b1": {"process_id": "b", "arguments": {"a": {"from_node": "a1"}}, "result": True},
+        }
+        # Original graphs are not mutated silently
+        assert a_flat == {
+            "a1": {"process_id": "a", "arguments": {}},
+        }
+        assert b_flat == {
+            "a1": {"process_id": "a", "arguments": {}},
+            "b1": {"process_id": "b", "arguments": {"a": {"from_node": "a1"}}},
+        }
+
+
 def test_build_and_flatten_simple():
     node = PGNode("foo")
     assert node.flat_graph() == {"foo1": {"process_id": "foo", "arguments": {}, "result": True}}
@@ -421,14 +507,14 @@ class TestMultiResult:
     def test_simple(self):
         multi = MultiResult([PGNode("foo"), PGNode("bar")])
         assert multi.flat_graph() == {
-            "foo1": {"process_id": "foo", "arguments": {}, "result": True},
+            "foo1": {"process_id": "foo", "arguments": {}},
             "bar1": {"process_id": "bar", "arguments": {}, "result": True},
         }
 
     def test_simple_duplicates(self):
         multi = MultiResult([PGNode("foo"), PGNode("foo")])
         assert multi.flat_graph() == {
-            "foo1": {"process_id": "foo", "arguments": {}, "result": True},
+            "foo1": {"process_id": "foo", "arguments": {}},
             "foo2": {"process_id": "foo", "arguments": {}, "result": True},
         }
 
@@ -442,7 +528,6 @@ def test_multi_save_result_same_root(self):
             "saveresult1": {
                 "process_id": "save_result",
                 "arguments": {"data": {"from_node": "loadcollection1"}, "format": "GTiff", "options": {}},
-                "result": True,
             },
             "saveresult2": {
                 "process_id": "save_result",