MNT: Explicit option for savelayout in to_dot_graph

astropy · Aug 27, 2023 · 22057d3 · 22057d3
1 parent 026a140
commit 22057d3
Showing 1 changed file with 16 additions and 2 deletions.
diff --git a/astropy/coordinates/transformations.py b/astropy/coordinates/transformations.py
@@ -492,7 +492,7 @@ def to_dot_graph(
         savefn : None or str
             The file name to save this graph to or `None` to not save
             to a file.
-        savelayout : str
+        savelayout : {"plain", "dot", "neato", "fdp", "sfdp", "circo", "twopi", "nop", "nop2", "osage", "patchwork"}
             The graphviz program to use to layout the graph (see
             graphviz_ for details) or 'plain' to just save the DOT graph
             content. Ignored if ``savefn`` is `None`.
@@ -571,7 +571,19 @@ def to_dot_graph(
             if savelayout == "plain":
                 with open(savefn, "w") as f:
                     f.write(dotgraph)
-            else:
+            # Options from https://graphviz.org/docs/layouts/
+            elif savelayout in (
+                "dot",
+                "neato",
+                "fdp",
+                "sfdp",
+                "circo",
+                "twopi",
+                "nop",
+                "nop2",
+                "osage",
+                "patchwork",
+            ):
                 args = [savelayout]
                 if saveformat is not None:
                     args.append("-T" + saveformat)
@@ -587,6 +599,8 @@ def to_dot_graph(
 
                 with open(savefn, "w") as f:
                     f.write(stdout)
+            else:
+                raise NotImplementedError(f'savelayout="{savelayout}" is not supported')
 
         return dotgraph