Generalize pytest-mpl to support non-Matplotlib visual output

[astrofrog]

I am currently trying to use the Solara testing framework to try and do unit/regression testing of the rendering of Jupyter widgets, and I was wishing that I could just re-use all the pytest-mpl infrastructure to do figure hashes, galleries, etc. In a discussion with @Cadair we realised that it is possible, although a little ugly right now:

import ipywidgets as widgets
import matplotlib.pyplot as plt
import playwright.sync_api
import pytest
from IPython.display import display
from PIL import Image

@pytest.mark.mpl_image_compare(savefig_kwargs={"dpi": 100})
def test_scatter_solara(
    tmp_path,
    solara_test,
    page_session: playwright.sync_api.Page,
    assert_solara_snapshot,
):

    from glue_jupyter import jglue
    from IPython.display import display

    app = jglue()
    app.add_data(a={"x": [1, 2, 3], "y": [2, 3, 4]})
    scatter = app.scatter2d(show=False)
    display(scatter._layout)

    button_sel = page_session.locator("text=No selection")
    button_sel.wait_for()
    screenshot = page_session.screenshot()

    with open(tmp_path / "screenshot.png", "wb") as f:
        f.write(screenshot)

    image = Image.open(tmp_path / "screenshot.png")
    size_x, size_y = image.size

    fig = plt.figure(figsize=(size_x / 100, size_y / 100))
    ax = fig.add_axes([0, 0, 1, 1])
    ax.imshow(image)
    plt.axis("off")

    return fig

Basically one can in principle take any image output and show it in a Matplotlib figure. However this is a bit hacky, and what would be a lot nicer would be if we could generalize this plugin to allow for alternatives to what actually generates the images, with Matplotlib being just one (and given the plugin name the default) option.

To keep this simple, we could also simply allow the returned value from tests to be a PIL Image instance as this would already allow any arbitrary image generation in the tests.

Does anyone have any thoughts on this? Any reason to not at least allow PIL returns?

[Cadair]

The return value could be a png file already? I.e. a path or bytes? The plugin could then copy it to wherever it expects?

[astrofrog]

Yes sounds good

[ConorMacBride]

I haven't tested it, but as long as the object returned has a savefig(fname, *, **kwargs) method it should just work. So it might work if you create a new class which stores the screenshot upon initialisation and also implements savefig. (We should add tests for this sort of use case.) That said, I think PIL and pathlike returns would be useful to support out of the box.

[astrofrog]

Thanks! Indeed it works for now with:

class DummyFigure:

    def __init__(self, png_bytes):
        self._png_bytes = png_bytes

    def savefig(self, filename_or_fileobj, *args, **kwargs):
        if isinstance(filename_or_fileobj, str):
            with open(filename_or_fileobj, 'wb') as f:
                f.write(self._png_bytes)
        else:
            filename_or_fileobj.write(self._png_bytes)

though it would be nice to have 'official' support for this and avoid this workaround.

[MatthewFlamm]

I think collecting how other packages approach this, or collecting the needs of other projects, will be helpful to define what this should look like. We have a purpose built image comparison for PyVista at https://github.com/pyvista/pytest-pyvista.

In my opinion, a purposeful API would be more useful than hacking into the existing infrastructure. In pytest-pyvista we have tests with multiple image comparisons per test, and the need to turn off and on image comparisons within a test. For example a figure could be intermediately rendered several times, but only the final state is compared to the cached image. Another example is that PyVista can plot and render immediately without returning an object. It isn't immediately obvious to me that the current structure of pytest-mpl handles these complexities without workarounds. Other packages probably have other types of interesting needs here too. Having a more flexible structure that isn't tied to returning an object with a specific method would be very useful.

A suggestion is that pytest-pyvista uses a fixture for image comparison. This gives us a lot more power for adding in flexibility. A fixture can yield an object into the test that can be modified in a very flexible way. This also allows downstream packages to create their own fixtures that consume this fixture to make the image generation hooks. This is more pytest-onic in my opinion.

Without being very familiar with pytest-mpl, I'm not sure if this would be too much work to internalize for this package however.

[weiji14]

I haven't tested it, but as long as the object returned has a savefig(fname, *, **kwargs) method it should just work. So it might work if you create a new class which stores the screenshot upon initialisation and also implements savefig. (We should add tests for this sort of use case.) That said, I think PIL and pathlike returns would be useful to support out of the box.

Just gonna chime in to say that we've been using this exact hack in PyGMT since 2017 (see the current savefig method here, and an example test). PyGMT actually uses a totally different rendering method (ghostscript) compared to matplotlib, but the PNG comparisons have been working reliably for 6 years+ now.

That said, maybe if pytest-mpl were able to use __repr_png__'s generated output to compare against a baseline PNG, that might be easier than needing downstream libraries to implement savefig? I see that Pillow has implemented __repr_png__ at https://github.com/python-pillow/Pillow/pull/1091, and Matplotlib/PyGMT has it too, so it could be a good standard to use.