ENH: make retreive substructure work also for referenced records

src/caosdb/utils/plantuml.py

@@ -34,10 +34,13 @@ plantuml FILENAME.pu -> FILENAME.png
 """
 
 import os
+import shutil
 
 import caosdb as db
 from caosdb.common.datatype import is_reference, get_referenced_recordtype
 
+from typing import Optional
+
 REFERENCE = "REFERENCE"
 
 
@@ -263,6 +266,15 @@ def retrieve_substructure(start_record_types, depth, result_id_set=None, result_
             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 prop.id not in result_id_set:
                 result_container.append(prop)
@@ -281,7 +293,10 @@ def retrieve_substructure(start_record_types, depth, result_id_set=None, result_
     return None
 
 
-def to_graphics(recordtypes, filename):
+def to_graphics(recordtypes: list[db.Entity], filename: str,
+                output_dirname: Optional[str] = None,
+                formats: list[str]=["tsvg"],
+                silent:bool=True):
     """Calls recordtypes_to_plantuml_string(), saves result to file and
     creates an svg image
 
@@ -293,17 +308,44 @@ def to_graphics(recordtypes, filename):
                   Iterable with the entities to be displayed.
     filename : str
                filename of the image without the extension(e.g. data_structure;
+               also without the preceeding path.
                data_structure.pu and data_structure.svg will be created.)
+    output_dirname : str
+                     the destination directory for the resulting images as defined by the "-o"
+                     option by plantuml
+                     default is to use current working dir
+    formats : list[str]
+              list of target formats as defined by the -t"..." options by plantuml, e.g. "tsvg"
+    silent : bool
+             Don't output messages.
     """
     pu = recordtypes_to_plantuml_string(recordtypes)
 
-    pu_filename = filename+".pu"
-    with open(pu_filename, "w") as pu_file:
-        pu_file.write(pu)
-
-    cmd = "plantuml -tsvg %s" % pu_filename
-    print("Executing:", cmd)
-
-    if os.system(cmd) != 0:
-        raise Exception("An error occured during the execution of plantuml. "
-                        "Is plantuml installed?")
+    if output_dirname is None:
+        output_dirname = os.getcwd()
+        
+    allowed_formats = [
+        "tpng", "tsvg", "teps", "tpdf", "tvdx", "txmi",
+        "tscxml", "thtml", "ttxt", "tutxt", "tlatex", "tlatex:nopreamble"]
+
+    with tempfile.TemporaryDirectory() as td:
+        
+        pu_filename = os.path.join(td, filename + ".pu")
+        with open(pu_filename, "w") as pu_file:
+            pu_file.write(pu)
+
+        for format in formats:
+            extension = format[1:]
+            if ":" in extension:
+                extension = extension[:extension.index(":")]
+                
+            if format not in allowed_formats:
+                raise RuntimeError("Format not allowed.")
+            cmd = "plantuml -{} {}".format(format, pu_filename)
+            if not silent:
+                print("Executing:", cmd)
+
+            if os.system(cmd) != 0:  # TODO: replace with subprocess.run
+                raise Exception("An error occured during the execution of plantuml. "
+                                "Is plantuml installed?")
+            shutil.copy(os.path.join(td, filename + "." + extension), output_dirname)