simplify display of matplotlib figures

Author: Christopher Doris

src/compat/matplotlib.jl

@@ -1,47 +1,108 @@
+struct PyPlotFigure
+    py::Py
+end
+export PyPlotFigure
+
+ispy(::PyPlotFigure) = true
+getpy(fig::PyPlotFigure) = fig.py
+
+Base.show(io::IO, mime::MIME"image/png", fig::PyPlotFigure) = _pyplot_show(io, mime, fig, "png")
+Base.show(io::IO, mime::MIME"image/jpeg", fig::PyPlotFigure) = _pyplot_show(io, mime, fig, "jpeg")
+Base.show(io::IO, mime::MIME"image/tiff", fig::PyPlotFigure) = _pyplot_show(io, mime, fig, "tiff")
+Base.show(io::IO, mime::MIME"image/svg+xml", fig::PyPlotFigure) = _pyplot_show(io, mime, fig, "svg")
+Base.show(io::IO, mime::MIME"application/pdf", fig::PyPlotFigure) = _pyplot_show(io, mime, fig, "pdf")
+
+Base.showable(::MIME"image/png", fig::PyPlotFigure) = _pyplot_showable(fig, "png")
+Base.showable(::MIME"image/jpeg", fig::PyPlotFigure) = _pyplot_showable(fig, "jpeg")
+Base.showable(::MIME"image/tiff", fig::PyPlotFigure) = _pyplot_showable(fig, "tiff")
+Base.showable(::MIME"image/svg+xml", fig::PyPlotFigure) = _pyplot_showable(fig, "svg")
+Base.showable(::MIME"application/pdf", fig::PyPlotFigure) = _pyplot_showable(fig, "pdf")
+
+function _pyplot_bytes(fig::Py, format::String)
+    buf = pyimport("io").BytesIO()
+    fig.savefig(buf, format=format)
+    return pyconvert(Vector{UInt8}, buf.getvalue())
+end
+
+function _pyplot_show(io::IO, mime::MIME, fig::PyPlotFigure, format::String)
+    try
+        write(io, _pyplot_bytes(fig.py, format))
+        return
+    catch exc
+        if exc isa PyException
+            throw(MethodError(show, (io, mime, fig)))
+        else
+            rethrow()
+        end
+    end
+end
+
+function _pyplot_showable(fig::PyPlotFigure, format::String)
+    try
+        _pyplot_bytes(fig.py, format)
+        return true
+    catch exc
+        if exc isa PyException
+            return false
+        else
+            rethrow()
+        end
+    end
+end
+
+"""
+    pyplot_gcf(; close=true)
+
+Get the current matplotlib/pyplot/seaborn/etc figure as an object displayable by Julia's display mechanism.
+
+If `close` is true, the figure is also closed.
+"""
+function pyplot_gcf(; close::Bool=true)
+    plt = pyimport("matplotlib.pyplot")
+    fig = plt.gcf()
+    close && plt.close(fig)
+    PyPlotFigure(fig)
+end
+export pyplot_gcf
+
 """
     pyplot_show([fig]; close=true, [format])
 
-Show the matplotlib/pyplot/seaborn/etc figure `fig`, or all open figures if not given, using Julia's display mechanism.
+Show the matplotlib/pyplot/seaborn/etc figure `fig` using Julia's display mechanism.
+
+If `fig` is not given, then all open figures are shown.
 
 If `close` is true, the figure is also closed.
 
 The `format` specifies the file format of the generated image.
-By default this is `pyplot.rcParams["savefig.format"]` or `"png"`.
-It can be one of `"png"`, `"jpg"`, `"jpeg"`, `"tif"`, `"tiff"`, `"svg"` or `"pdf"`.
+If not given, it is decided by the display.
 """
-function pyplot_show(fig; close::Bool = true, format::String = "")
+function pyplot_show(fig; close::Bool = true, format::Union{String,Nothing} = nothing)
     plt = pyimport("matplotlib.pyplot")
-    io = pyimport("io")
+    fig = Py(fig)
     if !pyisinstance(fig, plt.Figure)
         fig = plt.figure(fig)
     end
-    buf = io.BytesIO()
-    if isempty(format)
-        format = pyconvert(String, plt.rcParams.get("savefig.format", "png"))
-    end
-    if format ∉ ("png", "jpg", "jpeg", "tif", "tiff", "svg", "pdf")
-        error("invalid format: $format")
-    end
-    fig.savefig(buf, format=format)
-    data = pyconvert(Vector{UInt8}, buf.getvalue())
-    if close
-        plt.close(fig)
-    end
-    if format == "png"
-        display(MIME("image/png"), data)
+    fig = PyPlotFigure(fig)
+    if format === nothing
+        display(fig)
+    elseif format == "png"
+        display(MIME("image/png"), fig)
     elseif format in ("jpg", "jpeg")
-        display(MIME("image/jpeg"), data)
+        display(MIME("image/jpeg"), fig)
     elseif format in ("tif", "tiff")
-        display(MIME("image/tiff"), data)
+        display(MIME("image/tiff"), fig)
     elseif format == "svg"
-        display(MIME("image/svg+xml"), String(data))
+        display(MIME("image/svg+xml"), fig)
     elseif format == "pdf"
-        display(MIME("application/pdf"), data)
+        display(MIME("application/pdf"), fig)
     else
-        @assert false
+        error("invalid format: $format")
     end
-    nothing
+    close && plt.close(fig)
+    return
 end
+
 function pyplot_show(; opts...)
     plt = pysysmodule.modules.get("matplotlib.pyplot", nothing)
     if !pyisnone(plt)

src/compat/multimedia.jl

@@ -113,5 +113,5 @@ end
 function init_pyshow()
     pyshow_add_rule(pyshow_rule_mimebundle)
     pyshow_add_rule(pyshow_rule_repr)
-    pyshow_add_rule(pyshow_rule_savefig)
+    # pyshow_add_rule(pyshow_rule_savefig)
 end