Allow show to write the screenshot to a file-like object

cadquery/vis.py

@@ -13,7 +13,7 @@
 from .occ_impl.assembly import _loc2vtk, toVTKAssy
 from .occ_impl.nurbs import Curve, Surface
 
-from typing import Union, List, Tuple, Iterable, cast, Optional
+from typing import Union, List, Tuple, Iterable, cast, Optional, IO
 
 from OCP.TopoDS import TopoDS_Shape
 from OCP.Geom import Geom_BSplineSurface
@@ -402,7 +402,7 @@ def show(
     edges: bool = False,
     specular: bool = True,
     title: str = "CQ viewer",
-    screenshot: Optional[str] = None,
+    screenshot: Optional[Union[str, IO[bytes]]] = None,
     interact: bool = True,
     zoom: float = 1.0,
     roll: float = -35,
@@ -572,9 +572,15 @@ def show(
         win2image.Update()
 
         writer = vtkPNGWriter()
-        writer.SetFileName(screenshot)
         writer.SetInputConnection(win2image.GetOutputPort())
-        writer.Write()
+
+        if hasattr(screenshot, "write"):
+            writer.WriteToMemoryOn()
+            writer.Write()
+            screenshot.write(memoryview(writer.GetResult()).tobytes())
+        else:
+            writer.SetFileName(screenshot)
+            writer.Write()
 
     # start interaction
     if interact:

doc/vis.rst

@@ -92,6 +92,16 @@ camera position and windows size.
 
    show(b, width=800, height=800, screenshot='img.png', zoom=2, roll=-20, elevation=-30, interact=False)
 
+``screenshot`` also accepts a binary file-like object, for example to keep the image in memory.
+
+.. code-block:: python
+
+   from io import BytesIO
+
+   buf = BytesIO()
+   show(b, screenshot=buf, interact=False)
+   png = buf.getvalue()
+
 
 .. warning::
     Intermittent issues were observed with this functionality, please submit detailed bug reports in case

tests/test_vis.py

@@ -17,9 +17,11 @@
 )
 from vtkmodules.vtkRenderingAnnotation import vtkAnnotatedCubeActor
 from vtkmodules.vtkIOImage import vtkPNGWriter
+from vtkmodules.vtkImagingSources import vtkImageCanvasSource2D
 
 from pytest import fixture, raises, mark
 import os
+from io import BytesIO
 from pathlib import Path
 
 from typing import List
@@ -158,6 +160,39 @@ def test_screenshot(wp, patch_vtk):
     show(wp, interact=False, screenshot="img.png", trihedron=False, gradient=False)
 
 
+class FakeImage(FakeWin2Img):
+    """
+    Small solid image in place of the render window contents.
+    """
+
+    def __init__(self):
+
+        self.source = vtkImageCanvasSource2D()
+        self.source.SetExtent(0, 9, 0, 4, 0, 0)
+        self.source.SetScalarTypeToUnsignedChar()
+        self.source.SetNumberOfScalarComponents(3)
+        self.source.FillBox(0, 9, 0, 4)
+
+    def GetOutputPort(self, *args):
+
+        return self.source.GetOutputPort()
+
+
+def test_screenshot_to_memory(wp, patch_vtk, monkeypatch):
+
+    monkeypatch.setattr(vis, "vtkWindowToImageFilter", FakeImage)
+    monkeypatch.setattr(vis, "vtkPNGWriter", vtkPNGWriter)
+
+    buf = BytesIO()
+    show(wp, interact=False, screenshot=buf, trihedron=False, gradient=False)
+    png = buf.getvalue()
+
+    assert png.startswith(b"\x89PNG\r\n\x1a\n")
+    # width and height from the IHDR chunk
+    assert int.from_bytes(png[16:20], "big") == 10
+    assert int.from_bytes(png[20:24], "big") == 5
+
+
 @mark.skip(reason="running only smoke test for now")
 def test_screenshot_no_patch(wp, tmpdir, cwd):