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):