Merge 88e12f4dc2 into 0cd0d84016

diagrams/__init__.py

@@ -151,6 +151,9 @@ class Diagram:
             filename isn't given.
         :param filename: The output filename, without the extension (.png).
             If not given, it will be generated from the name.
+            If the environment variable DIAGRAMS_FILENAME_SUFFIX is set, its
+            value is appended to the filename (e.g. suffix ".generated"
+            turns "diagrams_image" into "diagrams_image.generated").
         :param direction: Data flow direction. Default is 'left to right'.
         :param curvestyle: Edge routing style. One of "ortho", "curved", "spline", or "polyline".
         :param outformat: Output file format. Default is 'png'.
@@ -179,6 +182,9 @@ class Diagram:
             filename = "diagrams_image"
         elif not filename:
             filename = "_".join(self.name.split()).lower()
+        filename_suffix = os.environ.get("DIAGRAMS_FILENAME_SUFFIX", "")
+        if filename_suffix:
+            filename = f"{filename}{filename_suffix}"
         self.filename = filename
         self.dot = Digraph(self.name, filename=self.filename, strict=strict)
 

tests/test_diagram.py

@@ -110,6 +110,21 @@ class DiagramTest(unittest.TestCase):
             Node("node1")
         self.assertTrue(os.path.exists(f"{self.name}.png"))
 
+    def test_filename_suffix_env_var(self):
+        self.name = "my_custom_name"
+        suffix = ".generated"
+        os.environ["DIAGRAMS_FILENAME_SUFFIX"] = suffix
+        try:
+            with Diagram(name="Example 1", filename=self.name, show=False):
+                Node("node1")
+            self.assertTrue(os.path.exists(f"{self.name}{suffix}.png"))
+        finally:
+            del os.environ["DIAGRAMS_FILENAME_SUFFIX"]
+            try:
+                os.remove(f"{self.name}{suffix}.png")
+            except FileNotFoundError:
+                pass
+
     def test_empty_name(self):
         """Check that providing an empty name don't crash, but save in a diagrams_image.xxx file."""
         self.name = "diagrams_image"