ENH: refactored the plantuml module to include more customization options and use temp dirs

src/caosdb/utils/plantuml.py

@@ -34,10 +34,15 @@ 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
+
+import tempfile
+
 REFERENCE = "REFERENCE"
 
 
@@ -79,13 +84,23 @@ class Grouped(object):
         return self.parents
 
 
-def recordtypes_to_plantuml_string(iterable):
+def recordtypes_to_plantuml_string(iterable,
+                                   add_properties: bool = True,
+                                   add_recordtypes: bool = True,
+                                   add_legend: bool = True,
+                                   style: str = "default"):
     """Converts RecordTypes into a string for PlantUML.
 
     This function obtains an iterable and returns a string which can
     be input into PlantUML for a representation of all RecordTypes in
     the iterable.
 
+    Current options for style
+    -------------------------
+
+    "default" - Standard rectangles with uml class circle and methods section
+    "salexan" - Round rectangles, hide circle and methods section
+
     Current limitations
     -------------------
 
@@ -140,10 +155,23 @@ def recordtypes_to_plantuml_string(iterable):
         return result
 
     result = "@startuml\n\n"
+
+    if style == "default":
         result += "skinparam classAttributeIconSize 0\n"
+    elif style == "salexan":
+        result += """skinparam roundcorner 20\n
+skinparam boxpadding 20\n
+\n
+hide methods\n
+hide circle\n
+"""
+    else:
+        raise ValueError("Unknown style.")
 
-    result += "package Properties #DDDDDD {\n"
     
+
+    if add_properties:
+        result += "package Properties #DDDDDD {\n"
         for p in properties:
             inheritances[p] = p.get_parents()
             dependencies[p] = []
@@ -163,6 +191,7 @@ def recordtypes_to_plantuml_string(iterable):
             result += "}\n\n"
         result += "}\n\n"
 
+    if add_recordtypes:
         result += "package RecordTypes #DDDDDD {\n"
 
         for c in classes:
@@ -200,6 +229,7 @@ def recordtypes_to_plantuml_string(iterable):
                 result += "\"{klass}\" *-- \"{dep}\"\n".format(
                     klass=c.name, dep=dep)
 
+    if add_legend:
         result += """
 
 package \"B is a subtype of A\" <<Rectangle>> {
@@ -264,6 +294,15 @@ def retrieve_substructure(start_record_types, depth, result_id_set=None, result_
                 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)
                 result_id_set.add(prop.id)
@@ -281,7 +320,14 @@ 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,
+                add_properties: bool = True,
+                add_recordtypes: bool = True,
+                add_legend: bool = True,
+                style: str = "default"):
     """Calls recordtypes_to_plantuml_string(), saves result to file and
     creates an svg image
 
@@ -293,17 +339,48 @@ 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 = recordtypes_to_plantuml_string(recordtypes,
+                                        add_properties,
+                                        add_recordtypes,
+                                        add_legend,
+                                        style)
+
+    if output_dirname is None:
+        output_dirname = os.getcwd()
         
-    pu_filename = filename+".pu"
+    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)
 
-    cmd = "plantuml -tsvg %s" % pu_filename
+        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:
+            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)