ENH: new function that recursively retrieves connected entities

src/caosdb/utils/plantuml.py

@@ -222,6 +222,63 @@ def to_graphics(recordtypes, filename):
     data_structure.svg will be created.
     """
     pu = recordtypes_to_plantuml_string(recordtypes)
+def retrieve_substructure(start_record_types, depth, result_id_set=None, result_container=None, cleanup=True):
+    """Recursively retrieves CaosDB record types and properties, starting
+    from given initial types up to a specific depth.
+
+    Parameters
+    ----------
+    start_record_types : Iterable[db.Entity]
+                         Iterable with the entities to be displayed. Starting from these
+                         entities more entities will be retrieved.
+    depth : int
+            The maximum depth up to which to retriev sub entities.
+    result_id_set : set[int]
+                    Used by recursion. Filled with already visited ids.
+    result_container : db.Container
+                       Used by recursion. Filled with already visited entities.
+    cleanup : bool
+              Used by recursion. If True return the resulting result_container.
+              Don't return anything otherwise.
+
+    Returns
+    -------
+    db.Container
+                A container containing all the retrieved entites or None if cleanup is False.
+    """
+    # Initialize the id set and result container for level zero recursion depth:
+    if result_id_set is None:
+        result_id_set = set()
+    if result_container is None:
+        result_container = db.Container()
+
+    for entity in start_record_types:
+        entity.retrieve()
+        if entity.id not in result_id_set:
+            result_container.append(entity)
+        result_id_set.add(entity.id)
+        for prop in entity.properties:
+            if is_reference(prop.datatype) and prop.datatype != db.FILE and depth > 0:
+                rt = db.RecordType(name=get_referenced_recordtype(prop.datatype)).retrieve()
+                retrieve_substructure([rt], depth-1, result_id_set, result_container, False)
+            else:
+                if prop.id not in result_id_set:
+                    result_container.append(prop)
+                result_id_set.add(prop.id)
+                
+        for parent in entity.parents:
+            rt = db.RecordType(id=parent.id).retrieve()
+            if parent.id not in result_id_set:
+                result_container.append(rt)
+            result_id_set.add(parent.id)
+            if depth > 0:
+                retrieve_substructure([rt], depth-1, result_id_set, result_container, False)
+
+    if cleanup:
+        return result_container
+    return None
+        
+    
 
     pu_filename = filename+".pu"
     with open(pu_filename, "w") as pu_file: