fix: properly handle None dependency_graph during merge (#1657)

* fix: properly handle None dependency_graph during merge

* fix: raise error for one None one not

* fix: change merge code to work when one graph exists and the other doesn't

* chore: refactor to simplify __add__ function

* test: ensure link occurred in process merge

* chore: fix dependencies due to new adsm

Author: dbirman

pyproject.toml

@@ -14,7 +14,7 @@ readme = "README.md"
 dynamic = ["version"]
 
 dependencies = [
-    'aind-data-schema-models>=4.2.11',
+    'aind-data-schema-models>=4.2.11,<5',
     'pydantic>=2.7, <2.12',
     'semver'
 ]

src/aind_data_schema/core/processing.py

@@ -12,7 +12,7 @@
 from aind_data_schema.base import AwareDatetimeWithDefault, DataCoreModel, DataModel, GenericModel
 from aind_data_schema.components.identifiers import Code
 from aind_data_schema.components.wrappers import AssetPath
-from aind_data_schema.utils.merge import merge_notes, merge_optional_list
+from aind_data_schema.utils.merge import merge_notes, merge_optional_list, merge_process_graph
 from aind_data_schema.utils.validators import TimeValidation
 
 
@@ -135,12 +135,14 @@ def rename_process(self, old_name: str, new_name: str) -> None:
                 break
         else:
             raise ValueError(f"Process '{old_name}' not found in data_processes.")
+
         # rename in dependency_graph
-        self.dependency_graph[new_name] = self.dependency_graph.pop(old_name)
-        # replace old_name in dependency_graph values
-        for value in self.dependency_graph.values():
-            if old_name in value:
-                value[value.index(old_name)] = new_name
+        if self.dependency_graph:
+            self.dependency_graph[new_name] = self.dependency_graph.pop(old_name)
+            # replace old_name in dependency_graph values
+            for value in self.dependency_graph.values():
+                if old_name in value:
+                    value[value.index(old_name)] = new_name
 
     @model_validator(mode="after")
     def order_processes(self) -> "Processing":
@@ -248,19 +250,19 @@ def __add__(self, other: "Processing") -> "Processing":
                     i += 1
                 other.rename_process(name, new_name)
 
-        # Merge process graphs - start with self's graph and update with other's graph
-        combined_process_graph = self.dependency_graph.copy()
-        combined_process_graph.update(other.dependency_graph)
+        merged_graph = merge_process_graph(
+            self.dependency_graph, other.dependency_graph, self.data_processes, other.data_processes
+        )
 
         # link self's output to other's input
         # note that this only makes sense if self has a single output process
         # and other has a single input process
-        if len(self.data_processes) > 0 and len(other.data_processes) > 0:
-            combined_process_graph[other.data_processes[0].name] = [self.data_processes[-1].name]
+        if merged_graph and len(self.data_processes) > 0 and len(other.data_processes) > 0:
+            merged_graph[other.data_processes[0].name] = [self.data_processes[-1].name]
 
         return Processing(
             pipelines=merge_optional_list(self.pipelines, other.pipelines),
             data_processes=self.data_processes + other.data_processes,
-            dependency_graph=combined_process_graph,
+            dependency_graph=merged_graph,
             notes=merge_notes(self.notes, other.notes),
         )

src/aind_data_schema/utils/merge.py

@@ -4,6 +4,34 @@
 from typing import Any, List, Optional
 
 
+def merge_process_graph(
+    graph1: Optional[dict],
+    graph2: Optional[dict],
+    processes1: List[Any],
+    processes2: List[Any],
+) -> Optional[dict]:
+    """Merge two process dependency graphs"""
+
+    # Merge process graphs - start with self's graph and update with other's graph
+    if graph1 and graph2:
+        merged_graph = graph1.copy()
+        merged_graph.update(graph2)
+    elif graph1 and not graph2:
+        merged_graph = graph1.copy()
+        # Add entries for other's processes
+        for process in processes2:
+            merged_graph[process.name] = []
+    elif graph2 and not graph1:
+        merged_graph = graph2.copy()
+        # Add entries for self's processes
+        for process in processes1:
+            merged_graph[process.name] = []
+    else:
+        merged_graph = None
+
+    return merged_graph
+
+
 def merge_str_tuple_lists(
     a: list[str | tuple[str, ...]], b: list[str | tuple[str, ...]]
 ) -> list[str | tuple[str, ...]]:

tests/test_composability_merge.py

@@ -311,6 +311,85 @@ def test_add_processing_objects(self):
         self.assertEqual(combined.data_processes[2].name, "Analysis_2")
         self.assertEqual(combined.data_processes[3].name, "Analysis_3")
 
+    def test_merge_dependency_graph(self):
+        """Test merging dependency graphs"""
+
+        t = datetime(2022, 11, 22, 8, 43, 00, tzinfo=timezone.utc)
+
+        # Test None/None case
+        p1 = Processing(
+            data_processes=[
+                DataProcess(
+                    experimenters=["Dr. Dan"],
+                    process_type=ProcessName.ANALYSIS,
+                    stage=ProcessStage.PROCESSING,
+                    start_date_time=t,
+                    code=Code(url="https://example.com", version="1.0"),
+                )
+            ],
+            dependency_graph=None,
+        )
+
+        p2 = Processing(
+            data_processes=[
+                DataProcess(
+                    experimenters=["Dr. Jane"],
+                    process_type=ProcessName.COMPRESSION,
+                    stage=ProcessStage.PROCESSING,
+                    start_date_time=t,
+                    code=Code(url="https://example.com", version="1.0"),
+                )
+            ],
+            dependency_graph=None,
+        )
+
+        combined = p1 + p2
+        self.assertIsNone(combined.dependency_graph)
+
+        # Test both set case
+        p3 = Processing(
+            data_processes=[
+                DataProcess(
+                    experimenters=["Dr. Dan"],
+                    process_type=ProcessName.ANALYSIS,
+                    stage=ProcessStage.PROCESSING,
+                    start_date_time=t,
+                    code=Code(url="https://example.com", version="1.0"),
+                )
+            ],
+            dependency_graph={"Analysis": []},
+        )
+
+        p4 = Processing(
+            data_processes=[
+                DataProcess(
+                    experimenters=["Dr. Jane"],
+                    process_type=ProcessName.COMPRESSION,
+                    stage=ProcessStage.PROCESSING,
+                    start_date_time=t,
+                    code=Code(url="https://example.com", version="1.0"),
+                )
+            ],
+            dependency_graph={"Compression": []},
+        )
+
+        combined = p3 + p4
+        self.assertIsNotNone(combined.dependency_graph)
+        self.assertEqual(len(combined.dependency_graph), 2)
+        self.assertIn("Analysis", combined.dependency_graph)
+        self.assertIn("Compression", combined.dependency_graph)
+        self.assertEqual(combined.dependency_graph["Compression"], ["Analysis"])
+
+        # Test self has graph, other doesn't
+        combined = p3 + p2
+        self.assertIsNotNone(combined.dependency_graph)
+        self.assertIn("Analysis", combined.dependency_graph)
+
+        # Test other has graph, self doesn't
+        combined = p2 + p4
+        self.assertIsNotNone(combined.dependency_graph)
+        self.assertIn("Compression", combined.dependency_graph)
+
 
 if __name__ == "__main__":
     unittest.main()

tests/test_utils_merge.py

@@ -1,13 +1,15 @@
 """ Tests for merge utilities """
 
 import unittest
+from unittest.mock import Mock
 
 from aind_data_schema.components.coordinates import CoordinateSystemLibrary
 from aind_data_schema.utils.merge import (
     merge_notes,
     merge_optional_list,
     merge_coordinate_systems,
     merge_str_tuple_lists,
+    merge_process_graph,
 )
 
 
@@ -193,5 +195,37 @@ def test_both_different(self):
             merge_coordinate_systems(self.CSA, self.CSB)
 
 
+class MergeProcessGraphTests(unittest.TestCase):
+    """Tests for merge_process_graph"""
+
+    def test_both_graphs_present(self):
+        """Test merging when both graphs are present"""
+        graph1 = {"proc1": ["proc2"], "proc2": []}
+        graph2 = {"proc3": ["proc4"], "proc4": []}
+        result = merge_process_graph(graph1, graph2, [], [])
+        self.assertEqual(result, {"proc1": ["proc2"], "proc2": [], "proc3": ["proc4"], "proc4": []})
+
+    def test_only_first_graph(self):
+        """Test when only first graph is present"""
+        graph1 = {"proc1": ["proc2"], "proc2": []}
+        proc = Mock()
+        proc.name = "proc3"
+        result = merge_process_graph(graph1, None, [], [proc])
+        self.assertEqual(result, {"proc1": ["proc2"], "proc2": [], "proc3": []})
+
+    def test_only_second_graph(self):
+        """Test when only second graph is present"""
+        graph2 = {"proc3": ["proc4"], "proc4": []}
+        proc = Mock()
+        proc.name = "proc1"
+        result = merge_process_graph(None, graph2, [proc], [])
+        self.assertEqual(result, {"proc3": ["proc4"], "proc4": [], "proc1": []})
+
+    def test_both_graphs_none(self):
+        """Test when both graphs are None"""
+        result = merge_process_graph(None, None, [], [])
+        self.assertIsNone(result)
+
+
 if __name__ == "__main__":
     unittest.main()