✨ Recursively query parents and children

lamindb/_parents.py

@@ -15,13 +15,39 @@
 if TYPE_CHECKING:
     from lnschema_core.types import StrField
 
+    from lamindb.core import QuerySet
+
 LAMIN_GREEN_LIGHTER = "#10b981"
 LAMIN_GREEN_DARKER = "#065f46"
 GREEN_FILL = "honeydew"
 TRANSFORM_EMOJIS = {"notebook": "📔", "app": "🖥️", "pipeline": "🧩"}
 is_run_from_ipython = getattr(builtins, "__IPYTHON__", False)
 
 
+# this is optimized to have fewer recursive calls
+# also len of QuerySet can be costly at times
+def _query_relatives(
+    records: QuerySet | list[Record],
+    kind: Literal["parents", "children"],
+    cls: type[HasParents],
+) -> QuerySet:
+    relatives = cls.objects.none()
+    if len(records) == 0:
+        return relatives
+    for record in records:
+        relatives = relatives.union(getattr(record, kind).all())
+    relatives = relatives.union(_query_relatives(relatives, kind, cls))
+    return relatives
+
+
+def query_parents(self) -> QuerySet:
+    return _query_relatives([self], "parents", self.__class__)
+
+
+def query_children(self) -> QuerySet:
+    return _query_relatives([self], "children", self.__class__)
+
+
 def _transform_emoji(transform: Transform):
     if transform is not None:
         return TRANSFORM_EMOJIS.get(transform.type, "💫")
@@ -474,9 +500,7 @@ def _df_edges_from_runs(df_values: list):
     return df
 
 
-METHOD_NAMES = [
-    "view_parents",
-]
+METHOD_NAMES = ["view_parents", "query_parents", "query_children"]
 
 if ln_setup._TESTING:  # type: ignore
     from inspect import signature

tests/core/test_parents.py

@@ -13,6 +13,26 @@ def test_view_parents():
     label2.delete()
 
 
+def test_query_parents_children():
+    label1 = ln.ULabel(name="label1")
+    label2 = ln.ULabel(name="label2")
+    label3 = ln.ULabel(name="label3")
+    label1.save()
+    label2.save()
+    label3.save()
+    label1.children.add(label2)
+    label2.children.add(label3)
+    parents = label3.query_parents()
+    assert len(parents) == 2
+    assert label1 in parents and label2 in parents
+    children = label1.query_children()
+    assert len(children) == 2
+    assert label2 in children and label3 in children
+    label1.delete()
+    label2.delete()
+    label3.delete()
+
+
 def test_add_emoji():
     record = ln.Transform(type="app")
     assert _add_emoji(record, label="transform") == "🖥️ transform"