diff --git a/src/caosdb/utils/plantuml.py b/src/caosdb/utils/plantuml.py
index c226a2fe1850f2c95b0a0ccef230842e9644c5e5..ca576ed1dd5072159dfb2abae937e4c69311715e 100644
--- a/src/caosdb/utils/plantuml.py
+++ b/src/caosdb/utils/plantuml.py
@@ -276,7 +276,8 @@ def retrieve_substructure(start_record_types, depth, result_id_set=None, result_
     Returns
     -------
     db.Container
-                A container containing all the retrieved entites or None if cleanup is False.
+                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:
@@ -290,18 +291,11 @@ def retrieve_substructure(start_record_types, depth, result_id_set=None, result_
             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)
-
-            # TODO: clean up this hack
-            # TODO: make it also work for files
-            if is_reference(prop.datatype) and prop.value is not None:
-                r = db.Record(id=prop.value).retrieve()
-                retrieve_substructure([r], depth-1, result_id_set, result_container, False)
-                if r.id not in result_id_set:
-                    result_container.append(r)
-                    result_id_set.add(r.id)
+            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)
 
             if prop.id not in result_id_set:
                 result_container.append(prop)
@@ -313,7 +307,8 @@ def retrieve_substructure(start_record_types, depth, result_id_set=None, result_
                 result_container.append(rt)
             result_id_set.add(parent.id)
             if depth > 0:
-                retrieve_substructure([rt], depth-1, result_id_set, result_container, False)
+                retrieve_substructure([rt], depth-1, result_id_set,
+                                      result_container, False)
 
     if cleanup:
         return result_container