Return filename from save_figure (matplotlib#27766)

Zybulon · tacaswell · web-flow · commit ff262a15e399 · 2024-10-28T22:36:52.000+01:00
---------

Co-authored-by: Thomas A Caswell &lt;tcaswell@gmail.com&gt;
diff --git a/doc/api/next_api_changes/behavior/27744-FM.rst b/doc/api/next_api_changes/behavior/27744-FM.rst
@@ -0,0 +1,10 @@
+``NavigationToolbar2.save_figure`` now returns filepath of saved figure
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+``NavigationToolbar2.save_figure`` function may return the filename of the saved figure.
+
+If a backend implements this functionality it should return `None`
+in the case where no figure is actually saved (because the user closed the dialog without saving).
+
+If the backend does not or can not implement this functionality (currently the Gtk4 backends
+and webagg backends do not) this method will return ``NavigationToolbar2.UNKNOWN_SAVED_STATUS``.
diff --git a/lib/matplotlib/backend_bases.py b/lib/matplotlib/backend_bases.py
@@ -2780,6 +2780,8 @@ class NavigationToolbar2:
         ('Save', 'Save the figure', 'filesave', 'save_figure'),
       )
 
+    UNKNOWN_SAVED_STATUS = object()
+
     def __init__(self, canvas):
         self.canvas = canvas
         canvas.toolbar = self
@@ -3194,7 +3196,26 @@ def on_tool_fig_close(e):
         return self.subplot_tool
 
     def save_figure(self, *args):
-        """Save the current figure."""
+        """
+        Save the current figure.
+
+        Backend implementations may choose to return
+        the absolute path of the saved file, if any, as
+        a string.
+
+        If no file is created then `None` is returned.
+
+        If the backend does not implement this functionality
+        then `NavigationToolbar2.UNKNOWN_SAVED_STATUS` is returned.
+
+        Returns
+        -------
+        str or `NavigationToolbar2.UNKNOWN_SAVED_STATUS` or `None`
+            The filepath of the saved figure.
+            Returns `None` if figure is not saved.
+            Returns `NavigationToolbar2.UNKNOWN_SAVED_STATUS` when
+            the backend does not provide the information.
+        """
         raise NotImplementedError
 
     def update(self):
diff --git a/lib/matplotlib/backend_bases.pyi b/lib/matplotlib/backend_bases.pyi
@@ -402,6 +402,7 @@ class _Mode(str, Enum):
 
 class NavigationToolbar2:
     toolitems: tuple[tuple[str, ...] | tuple[None, ...], ...]
+    UNKNOWN_SAVED_STATUS: object
     canvas: FigureCanvasBase
     mode: _Mode
     def __init__(self, canvas: FigureCanvasBase) -> None: ...
@@ -437,7 +438,7 @@ class NavigationToolbar2:
     def push_current(self) -> None: ...
     subplot_tool: widgets.SubplotTool
     def configure_subplots(self, *args): ...
-    def save_figure(self, *args) -> None: ...
+    def save_figure(self, *args) -> str | None | object: ...
     def update(self) -> None: ...
     def set_history_buttons(self) -> None: ...
 
diff --git a/lib/matplotlib/backends/_backend_tk.py b/lib/matplotlib/backends/_backend_tk.py
@@ -870,7 +870,7 @@ def save_figure(self, *args):
             )
 
         if fname in ["", ()]:
-            return
+            return None
         # Save dir for next time, unless empty str (i.e., use cwd).
         if initialdir != "":
             mpl.rcParams['savefig.directory'] = (
@@ -885,6 +885,7 @@ def save_figure(self, *args):
 
         try:
             self.canvas.figure.savefig(fname, format=extension)
+            return fname
         except Exception as e:
             tkinter.messagebox.showerror("Error saving file", str(e))
 
diff --git a/lib/matplotlib/backends/backend_gtk3.py b/lib/matplotlib/backends/backend_gtk3.py
@@ -339,7 +339,7 @@ def __init__(self, canvas):
     def save_figure(self, *args):
         dialog = Gtk.FileChooserDialog(
             title="Save the figure",
-            parent=self.canvas.get_toplevel(),
+            transient_for=self.canvas.get_toplevel(),
             action=Gtk.FileChooserAction.SAVE,
             buttons=(Gtk.STOCK_CANCEL, Gtk.ResponseType.CANCEL,
                      Gtk.STOCK_SAVE,   Gtk.ResponseType.OK),
@@ -371,16 +371,17 @@ def on_notify_filter(*args):
         fmt = self.canvas.get_supported_filetypes_grouped()[ff.get_name()][0]
         dialog.destroy()
         if response != Gtk.ResponseType.OK:
-            return
+            return None
         # Save dir for next time, unless empty str (which means use cwd).
         if mpl.rcParams['savefig.directory']:
             mpl.rcParams['savefig.directory'] = os.path.dirname(fname)
         try:
             self.canvas.figure.savefig(fname, format=fmt)
+            return fname
         except Exception as e:
             dialog = Gtk.MessageDialog(
-                parent=self.canvas.get_toplevel(), message_format=str(e),
-                type=Gtk.MessageType.ERROR, buttons=Gtk.ButtonsType.OK)
+                transient_for=self.canvas.get_toplevel(), text=str(e),
+                message_type=Gtk.MessageType.ERROR, buttons=Gtk.ButtonsType.OK)
             dialog.run()
             dialog.destroy()
 
diff --git a/lib/matplotlib/backends/backend_gtk4.py b/lib/matplotlib/backends/backend_gtk4.py
@@ -393,6 +393,7 @@ def on_response(dialog, response):
                 msg.show()
 
         dialog.show()
+        return self.UNKNOWN_SAVED_STATUS
 
 
 class ToolbarGTK4(ToolContainerBase, Gtk.Box):
diff --git a/lib/matplotlib/backends/backend_macosx.py b/lib/matplotlib/backends/backend_macosx.py
@@ -142,6 +142,7 @@ def save_figure(self, *args):
         if mpl.rcParams['savefig.directory']:
             mpl.rcParams['savefig.directory'] = os.path.dirname(filename)
         self.canvas.figure.savefig(filename)
+        return filename
 
 
 class FigureManagerMac(_macosx.FigureManager, FigureManagerBase):
diff --git a/lib/matplotlib/backends/backend_qt.py b/lib/matplotlib/backends/backend_qt.py
@@ -835,6 +835,7 @@ def save_figure(self, *args):
                     self, "Error saving file", str(e),
                     QtWidgets.QMessageBox.StandardButton.Ok,
                     QtWidgets.QMessageBox.StandardButton.NoButton)
+        return fname
 
     def set_history_buttons(self):
         can_backward = self._nav_stack._pos > 0
diff --git a/lib/matplotlib/backends/backend_webagg_core.py b/lib/matplotlib/backends/backend_webagg_core.py
@@ -403,8 +403,9 @@ def remove_rubberband(self):
         self.canvas.send_event("rubberband", x0=-1, y0=-1, x1=-1, y1=-1)
 
     def save_figure(self, *args):
-        """Save the current figure"""
+        """Save the current figure."""
         self.canvas.send_event('save')
+        return self.UNKNOWN_SAVED_STATUS
 
     def pan(self):
         super().pan()
diff --git a/lib/matplotlib/backends/backend_wx.py b/lib/matplotlib/backends/backend_wx.py
@@ -1143,6 +1143,7 @@ def save_figure(self, *args):
                 mpl.rcParams["savefig.directory"] = str(path.parent)
             try:
                 self.canvas.figure.savefig(path, format=fmt)
+                return path
             except Exception as e:
                 dialog = wx.MessageDialog(
                     parent=self.canvas.GetParent(), message=str(e),
diff --git a/lib/matplotlib/tests/test_backend_gtk3.py b/lib/matplotlib/tests/test_backend_gtk3.py
@@ -1,6 +1,8 @@
+import os
 from matplotlib import pyplot as plt
 
 import pytest
+from unittest import mock
 
 
 @pytest.mark.backend("gtk3agg", skip_on_importerror=True)
@@ -46,3 +48,27 @@ def receive(event):
     fig.canvas.mpl_connect("draw_event", send)
     fig.canvas.mpl_connect("key_press_event", receive)
     plt.show()
+
+
+@pytest.mark.backend("gtk3agg", skip_on_importerror=True)
+def test_save_figure_return():
+    from gi.repository import Gtk
+    fig, ax = plt.subplots()
+    ax.imshow([[1]])
+    with mock.patch("gi.repository.Gtk.FileFilter") as fileFilter:
+        filt = fileFilter.return_value
+        filt.get_name.return_value = "Portable Network Graphics"
+        with mock.patch("gi.repository.Gtk.FileChooserDialog") as dialogChooser:
+            dialog = dialogChooser.return_value
+            dialog.get_filter.return_value = filt
+            dialog.get_filename.return_value = "foobar.png"
+            dialog.run.return_value = Gtk.ResponseType.OK
+            fname = fig.canvas.manager.toolbar.save_figure()
+            os.remove("foobar.png")
+            assert fname == "foobar.png"
+
+            with mock.patch("gi.repository.Gtk.MessageDialog"):
+                dialog.get_filename.return_value = None
+                dialog.run.return_value = Gtk.ResponseType.OK
+                fname = fig.canvas.manager.toolbar.save_figure()
+                assert fname is None
diff --git a/lib/matplotlib/tests/test_backend_macosx.py b/lib/matplotlib/tests/test_backend_macosx.py
@@ -1,6 +1,7 @@
 import os
 
 import pytest
+from unittest import mock
 
 import matplotlib as mpl
 import matplotlib.pyplot as plt
@@ -50,3 +51,17 @@ def new_choose_save_file(title, directory, filename):
 def test_ipython():
     from matplotlib.testing import ipython_in_subprocess
     ipython_in_subprocess("osx", {(8, 24): "macosx", (7, 0): "MacOSX"})
+
+
+@pytest.mark.backend('macosx')
+def test_save_figure_return():
+    fig, ax = plt.subplots()
+    ax.imshow([[1]])
+    prop = "matplotlib.backends._macosx.choose_save_file"
+    with mock.patch(prop, return_value="foobar.png"):
+        fname = fig.canvas.manager.toolbar.save_figure()
+        os.remove("foobar.png")
+        assert fname == "foobar.png"
+    with mock.patch(prop, return_value=None):
+        fname = fig.canvas.manager.toolbar.save_figure()
+        assert fname is None
diff --git a/lib/matplotlib/tests/test_backend_qt.py b/lib/matplotlib/tests/test_backend_qt.py
@@ -226,6 +226,20 @@ def test_figureoptions():
         fig.canvas.manager.toolbar.edit_parameters()
 
 
+@pytest.mark.backend('QtAgg', skip_on_importerror=True)
+def test_save_figure_return():
+    fig, ax = plt.subplots()
+    ax.imshow([[1]])
+    prop = "matplotlib.backends.qt_compat.QtWidgets.QFileDialog.getSaveFileName"
+    with mock.patch(prop, return_value=("foobar.png", None)):
+        fname = fig.canvas.manager.toolbar.save_figure()
+        os.remove("foobar.png")
+        assert fname == "foobar.png"
+    with mock.patch(prop, return_value=(None, None)):
+        fname = fig.canvas.manager.toolbar.save_figure()
+        assert fname is None
+
+
 @pytest.mark.backend('QtAgg', skip_on_importerror=True)
 def test_figureoptions_with_datetime_axes():
     fig, ax = plt.subplots()
diff --git a/lib/matplotlib/tests/test_backend_tk.py b/lib/matplotlib/tests/test_backend_tk.py
@@ -196,6 +196,23 @@ class Toolbar(NavigationToolbar2Tk):
     print("success")
 
 
+@_isolated_tk_test(success_count=2)
+def test_save_figure_return():
+    import matplotlib.pyplot as plt
+    from unittest import mock
+    fig = plt.figure()
+    prop = "tkinter.filedialog.asksaveasfilename"
+    with mock.patch(prop, return_value="foobar.png"):
+        fname = fig.canvas.manager.toolbar.save_figure()
+        os.remove("foobar.png")
+        assert fname == "foobar.png"
+        print("success")
+    with mock.patch(prop, return_value=""):
+        fname = fig.canvas.manager.toolbar.save_figure()
+        assert fname is None
+        print("success")
+
+
 @_isolated_tk_test(success_count=1)
 def test_canvas_focus():
     import tkinter as tk
