Speed up Collection.add() and save() on chained nodes

pipeline/src/base.py

@@ -204,27 +204,36 @@ def is_valid(self):
         return len(failures) == 0
 
     @property
-    def links(self):
+    def _direct_links(self):
         """
-        Return a list of attributes that reference other metadata nodes
+        Return the linked nodes directly referenced by this node's attributes.
+
+        Embedded nodes are not listed, but the linked nodes inside them are.
         """
         _links = []
         for property in self.__class__.properties:
             value = getattr(self, property.name)
-            if property.multiple:
-                if not isinstance(value, (tuple, list)):
-                    value = [value]
-                for item in value:
-                    if isinstance(item, LinkedMetadata):
-                        _links.append(item)
-                    if hasattr(item, "links"):
-                        _links.extend(item.links)
-            elif isinstance(value, LinkedMetadata):
-                _links.append(value)
-            if hasattr(value, "links"):
-                _links.extend(value.links)
+            if not property.multiple or not isinstance(value, (tuple, list)):
+                value = [value]
+            for item in value:
+                if isinstance(item, LinkedMetadata):
+                    _links.append(item)
+                elif isinstance(item, EmbeddedMetadata):
+                    _links.extend(item._direct_links)
         return _links
 
+    @property
+    def links(self):
+        """
+        Return all linked nodes reachable from this node, at any depth.
+
+        Nodes are listed depth-first: each directly linked node is followed by
+        its own links. Embedded nodes are not listed, but the linked nodes inside
+        them are. Duplicates are kept: a node reachable through several paths
+        appears once per path.
+        """
+        return [node for child in self._direct_links for node in (child, *child.links)]
+
     def _resolve_links(self, node_lookup):
         """Replace `Link` attributes with typed Nodes where possible"""
         for property in self.__class__.properties:

pipeline/src/collection.py

@@ -55,7 +55,7 @@ def _add_node(self, node):
         if node.id is None:
             node.id = self._get_blank_node_identifier()
         self.nodes[node.id] = node
-        for linked_node in node.links:
+        for linked_node in node._direct_links:
             self._add_node(linked_node)
 
     def _get_blank_node_identifier(self):
@@ -133,7 +133,7 @@ def save(self, path, individual_files=False, include_empty_properties=False, gro
             else:
                 data_context = {"@vocab": "https://openminds.om-i.org/props/"}
 
-            for linked_node in node.links:
+            for linked_node in node._direct_links:
                 self._add_node(linked_node)
         # Now we can actually save the nodes
         if not individual_files:

pipeline/tests/test_collections.py

@@ -159,3 +159,45 @@ def test_collection_sort_by_id():
     }
 
     assert saved_data == expected_saved_data
+
+
+class CountingCollection(Collection):
+    def __init__(self):
+        self.calls = 0
+        super().__init__()
+
+    def _add_node(self, node):
+        self.calls += 1
+        super()._add_node(node)
+
+
+def build_chain(n):
+    """Return n subject states, each one descended from the previous one."""
+    states = [omcore.SubjectState(lookup_label="s0")]
+    for i in range(1, n):
+        states.append(omcore.SubjectState(lookup_label=f"s{i}", descended_from=states[-1]))
+    return states
+
+
+def test_add_visits_each_node_once():
+    # adding a chain must not revisit nodes already reached through the recursion
+    # (iterating over node.links instead of the direct links made 2**(n-1) calls here)
+    states = build_chain(12)
+    collection = CountingCollection()
+    collection.add(states[-1])
+    assert len(collection) == 12
+    assert collection.calls == 12
+
+
+def test_blank_node_ids_follow_depth_first_order():
+    _, _, third = build_chain(3)
+    other = omcore.SubjectState(lookup_label="other")
+    subject = omcore.Subject(lookup_label="S", studied_states=[third, other])
+    collection = Collection(subject)
+    assert [(node_id, node.lookup_label) for node_id, node in collection.nodes.items()] == [
+        ("_:000000", "S"),
+        ("_:000001", "s2"),
+        ("_:000002", "s1"),
+        ("_:000003", "s0"),
+        ("_:000004", "other"),
+    ]

pipeline/tests/test_links.py

@@ -0,0 +1,47 @@
+"""
+Tests of Node.links and Node._direct_links.
+"""
+
+from openminds.latest.controlled_terms.age_reference import AgeReference
+from openminds.latest.controlled_terms.species import Species
+from openminds.latest.controlled_terms.unit_of_measurement import UnitOfMeasurement
+from openminds.latest.core.miscellaneous.quantitative_value import QuantitativeValue
+from openminds.latest.core.research.specimen_age import SpecimenAge
+from openminds.latest.core.research.subject import Subject
+from openminds.latest.core.research.subject_state import SubjectState
+
+
+def build_state(label, weeks, previous=None):
+    return SubjectState(
+        lookup_label=label,
+        descended_from=previous,
+        age=SpecimenAge(
+            age=QuantitativeValue(value=weeks, unit=UnitOfMeasurement.week),
+            reference=AgeReference.birth,
+        ),
+    )
+
+
+def test_direct_links_returns_only_direct_links():
+    # only the first layer is returned
+    first = SubjectState(lookup_label="first")
+    second = SubjectState(lookup_label="second", descended_from=first)
+    third = SubjectState(lookup_label="third", descended_from=second)
+    assert third._direct_links == [second]
+
+
+def test_direct_links_look_through_embedded_nodes():
+    # embedded nodes (SpecimenAge, QuantitativeValue) are looked through
+    state = build_state("s", 1)
+    assert state._direct_links == [UnitOfMeasurement.week, AgeReference.birth]
+
+
+def test_links_order_and_duplicates():
+    # each directly linked node is followed by its own links (depth-first),
+    # embedded nodes are looked through, and duplicates are kept
+    week, birth = UnitOfMeasurement.week, AgeReference.birth
+    a = build_state("A", 1)
+    b = build_state("B", 2, previous=a)
+    subject = Subject(lookup_label="S", species=Species.mus_musculus, studied_states=[a, b])
+
+    assert subject.links == [Species.mus_musculus, a, week, birth, b, week, birth, a, week, birth]