Merge pull request #341 from brainglobe/screenshot-types

Added support for multiple screenshot file formats

brainrender/render.py

@@ -347,9 +347,11 @@ def screenshot(self, name=None, scale=None):
             self.render(interactive=False)
 
         timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
-        name = name or f"brainrender_screenshot_{timestamp}"
-        if ".png" not in name:
-            name += ".png"
+        name = Path(name or f"brainrender_screenshot_{timestamp}")
+
+        # If no suffix is provided or it an unsupported format, default to .png
+        if name.suffix not in [".png", ".eps", ".pdf", ".svg", ".jpg"]:
+            name = name.with_suffix(".png")
 
         scale = scale or settings.SCREENSHOT_SCALE
 

tests/test_scene.py

@@ -1,4 +1,5 @@
 import shutil
+from pathlib import Path
 
 import pytest
 
@@ -101,12 +102,34 @@ def test_scene_slice():
 
 
 @pytest.mark.parametrize(
-    "name, scale", [("test", 2), (None, None), (None, 1), ("test2", None)]
+    "name, scale, expected_suffix",
+    [
+        ("test", 2, ".png"),
+        (None, None, ".png"),
+        (None, 1, ".png"),
+        ("test2", None, ".png"),
+        ("test.jpg", 1, ".jpg"),
+        ("test.eps", 1, ".eps"),
+        ("test.svg", 1, ".svg"),
+        ("test.pdf", 1, ".pdf"),
+        ("test.tiff", 1, ".png"),
+    ],
 )
-def test_scene_screenshot(name, scale):
-    s = Scene(screenshots_folder="tests/screenshots")
-    s.screenshot(name=name, scale=scale)
-    shutil.rmtree("tests/screenshots")
+def test_scene_screenshot(name, scale, expected_suffix):
+    screenshot_folder = Path.home() / "test_screenshots"
+    s = Scene(screenshots_folder=screenshot_folder)
+    out_path = s.screenshot(name=name, scale=scale)
+
+    assert Path(out_path).suffix == expected_suffix
+
+    # Vedo exports eps and svg files as gzipped files
+    # Append the .gz suffix to the expected path to check if file exists
+    if expected_suffix in [".eps", ".svg"]:
+        out_path += ".gz"
+
+    assert Path(out_path).exists()
+
+    shutil.rmtree(screenshot_folder)
     del s