From 645a2af43c2ad2c135f7fabdbe986c3c9b7862e0 Mon Sep 17 00:00:00 2001
From: lucas-santos-f <162367088+lucas-santos-f@users.noreply.github.com>
Date: Tue, 10 Nov 2026 15:58:12 -0300
Subject: [PATCH] Allow show to write the screenshot to a file-like object

screenshot now also takes a binary file-like object such as io.BytesIO.
The PNG is then written to memory with vtkPNGWriter.WriteToMemoryOn and
copied into it, which avoids a temporary file when the image is only
needed in memory (docs, notebooks).

Fixes #1760
---
 cadquery/vis.py   | 14 ++++++++++----
 doc/vis.rst       | 10 ++++++++++
 tests/test_vis.py | 35 +++++++++++++++++++++++++++++++++++
 3 files changed, 55 insertions(+), 4 deletions(-)

diff --git a/cadquery/vis.py b/cadquery/vis.py
index d79278990..dcc619a09 100644
--- a/cadquery/vis.py
+++ b/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:
diff --git a/doc/vis.rst b/doc/vis.rst
index 44ffa7047..0ef2b6cf8 100644
--- a/doc/vis.rst
+++ b/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
diff --git a/tests/test_vis.py b/tests/test_vis.py
index b59367d73..b9908c4dd 100644
--- a/tests/test_vis.py
+++ b/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):