Re-fix renderers outliving underlying streams.

e.g. `figure(layout="constrained"); imshow([[0, 1]], rasterized=True); savefig("/tmp/test.pdf")`

Author: anntzer

ext/_mplcairo.cpp

@@ -350,7 +350,7 @@ GraphicsContextRenderer::GraphicsContextRenderer(
 {}
 
 cairo_t* GraphicsContextRenderer::cr_from_fileformat_args(
-  StreamSurfaceType type, std::optional<py::object> file,
+  StreamSurfaceType type, py::object file,
   double width, double height, double dpi)
 {
   auto surface_create_for_stream =
@@ -384,22 +384,17 @@ cairo_t* GraphicsContextRenderer::cr_from_fileformat_args(
       "cairo was built without {.name} support"_format(type)
       .cast<std::string>()};
   }
-  auto const& cb = file
-    ? cairo_write_func_t{
-      [](void* closure, unsigned char const* data, unsigned int length)
+  auto const& cb =
+    [](void* closure, unsigned char const* data, unsigned int length)
         -> cairo_status_t {
-          auto const& write =
-            py::reinterpret_borrow<py::object>(static_cast<PyObject*>(closure));
-          auto const& written =
-            write(py::memoryview::from_memory(data, length)).cast<unsigned int>();
-          return  // NOTE: This does not appear to affect the context status.
-            written == length ? CAIRO_STATUS_SUCCESS : CAIRO_STATUS_WRITE_ERROR;
-        }}
-    : nullptr;
-  py::object write =
-    // TODO: Why does py::none() not work here?
-    file ? file->attr("write") : py::reinterpret_borrow<py::object>(Py_None);
-
+      auto const& write =
+        py::reinterpret_borrow<py::object>(static_cast<PyObject*>(closure));
+      auto const& written =
+        write(py::memoryview::from_memory(data, length)).cast<unsigned int>();
+      return  // NOTE: This does not appear to affect the context status.
+        written == length ? CAIRO_STATUS_SUCCESS : CAIRO_STATUS_WRITE_ERROR;
+    };
+  auto const& write = file.attr("write");
   auto const& surface =
     surface_create_for_stream(cb, write.ptr(), width, height);
   cairo_surface_set_fallback_resolution(surface, dpi, dpi);
@@ -416,7 +411,7 @@ cairo_t* GraphicsContextRenderer::cr_from_fileformat_args(
 }
 
 GraphicsContextRenderer::GraphicsContextRenderer(
-  StreamSurfaceType type, std::optional<py::object> file,
+  StreamSurfaceType type, py::object file,
   double width, double height, double dpi) :
   GraphicsContextRenderer{
     cr_from_fileformat_args(type, file, width, height, dpi), width, height,
@@ -2137,8 +2132,7 @@ Only intended for debugging purposes.
     .def(py::init<double, double, double>())
     .def(py::init<
          py::object, double, double, double, std::tuple<double, double>>())
-    .def(py::init<
-         StreamSurfaceType, std::optional<py::object>, double, double, double>())
+    .def(py::init<StreamSurfaceType, py::object, double, double, double>())
     .def(
       py::pickle(
         [](GraphicsContextRenderer const& gcr) -> py::tuple {

ext/_mplcairo.h

@@ -63,10 +63,10 @@ class GraphicsContextRenderer {
     double width, double height, double dpi,
     std::tuple<double, double> device_scales);
   static cairo_t* cr_from_fileformat_args(
-    StreamSurfaceType type, std::optional<py::object> file,
+    StreamSurfaceType type, py::object file,
     double width, double height, double dpi);
   GraphicsContextRenderer(
-    StreamSurfaceType type, std::optional<py::object> file,
+    StreamSurfaceType type, py::object file,
     double width, double height, double dpi);
 
   static GraphicsContextRenderer make_pattern_gcr(cairo_surface_t* cr);

src/mplcairo/base.py

@@ -2,6 +2,7 @@
 import contextlib
 from functools import partial, partialmethod
 from gzip import GzipFile
+from io import BytesIO
 import logging
 import os
 from pathlib import Path
@@ -75,7 +76,7 @@ def from_pycairo_ctx(cls, ctx, width, height, dpi, orig_scale):
 
     @classmethod
     def _for_fmt_output(cls, fmt, stream, width, height, dpi):
-        if stream is not None and cbook.file_requires_unicode(stream):
+        if cbook.file_requires_unicode(stream):
             if fmt in [_StreamSurfaceType.PS, _StreamSurfaceType.EPS]:
                 # PS is (typically) ASCII -- Language Reference, section 3.2.
                 stream = _BytesWritingWrapper(stream, "ascii")
@@ -286,7 +287,8 @@ def _print_vector(self, renderer_factory,
                     self.figure.draw(renderer)
             except Exception as exc:
                 draw_raises_done = type(exc).__name__ == "Done"
-                raise
+                if not draw_raises_done:  # Else, will be re-raised below.
+                    raise
             finally:
                 # _finish() corresponds finalize() in Matplotlib's PDF and SVG
                 # backends; it is inlined in Matplotlib's PS backend.  It must
@@ -295,14 +297,14 @@ def _print_vector(self, renderer_factory,
                 # to be done before closing the stream.
                 renderer._finish()
         # ... but sometimes, Matplotlib *wants* a renderer to outlast the
-        # stream's lifetime, specifically to measure text extents.  This
-        # is done on Matplotlib's side by making Figure.draw throw a Done
-        # exception.  We catch that exception and swap in a no-write renderer
-        # (stream = None) in that case.
+        # stream's lifetime, specifically to measure text extents.  To do so,
+        # Matplotlib makes Figure.draw throw a Done exception.  We catch that
+        # exception and swap in a dummy renderer (stream = BytesIO()) in that
+        # case (no actual writing will be performed).
         if draw_raises_done:
-            renderer = renderer_factory(None, *self.figure.bbox.size, dpi)
+            renderer = renderer_factory(BytesIO(), *self.figure.bbox.size, dpi)
             with _LOCK:
-                self.figure.draw(renderer)
+                self.figure.draw(renderer)  # Should raise Done().
 
     print_pdf = partialmethod(
         _print_vector, GraphicsContextRendererCairo._for_pdf_output)