diff --git a/src/caosdb/utils/plantuml.py b/src/caosdb/utils/plantuml.py index 75b96ee3aa28ba916adc69e418de398abfe23356..31b562f983d5a189a352ac33758182269dc6dc40 100644 --- a/src/caosdb/utils/plantuml.py +++ b/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: