Merge pull request #22180 from greglucas/macosx-savefig

jklymak · web-flow · commit dd001edbbacc · 2022-05-18T09:02:34.000+02:00
ENH: Use rcParams savefig.directory on macosx backend
diff --git a/lib/matplotlib/backends/backend_macosx.py b/lib/matplotlib/backends/backend_macosx.py
@@ -1,3 +1,5 @@
+import os
+
 import matplotlib as mpl
 from matplotlib import _api, cbook
 from matplotlib._pylab_helpers import Gcf
@@ -115,10 +117,15 @@ def remove_rubberband(self):
         self.canvas.remove_rubberband()
 
     def save_figure(self, *args):
+        directory = os.path.expanduser(mpl.rcParams['savefig.directory'])
         filename = _macosx.choose_save_file('Save the figure',
+                                            directory,
                                             self.canvas.get_default_filename())
         if filename is None:  # Cancel
             return
+        # Save dir for next time, unless empty str (which means use cwd).
+        if mpl.rcParams['savefig.directory']:
+            mpl.rcParams['savefig.directory'] = os.path.dirname(filename)
         self.canvas.figure.savefig(filename)
 
     def prepare_configure_subplots(self):
diff --git a/lib/matplotlib/tests/test_backend_macosx.py b/lib/matplotlib/tests/test_backend_macosx.py
@@ -1,10 +1,13 @@
+import os
+
 import pytest
 
+import matplotlib as mpl
 import matplotlib.pyplot as plt
-
-
-pytest.importorskip("matplotlib.backends.backend_macosx",
-                    reason="These are mac only tests")
+try:
+    from matplotlib.backends import _macosx
+except ImportError:
+    pytest.skip("These are mac only tests", allow_module_level=True)
 
 
 @pytest.mark.backend('macosx')
@@ -18,3 +21,26 @@ def test_cached_renderer():
     fig = plt.figure(2)
     fig.draw_without_rendering()
     assert fig._cachedRenderer is not None
+
+
+@pytest.mark.backend('macosx')
+def test_savefig_rcparam(monkeypatch, tmp_path):
+
+    def new_choose_save_file(title, directory, filename):
+        # Replacement function instead of opening a GUI window
+        # Make a new directory for testing the update of the rcParams
+        assert directory == str(tmp_path)
+        os.makedirs(f"{directory}/test")
+        return f"{directory}/test/{filename}"
+
+    monkeypatch.setattr(_macosx, "choose_save_file", new_choose_save_file)
+    fig = plt.figure()
+    with mpl.rc_context({"savefig.directory": tmp_path}):
+        fig.canvas.toolbar.save_figure()
+        # Check the saved location got created
+        save_file = f"{tmp_path}/test/{fig.canvas.get_default_filename()}"
+        assert os.path.exists(save_file)
+
+        # Check the savefig.directory rcParam got updated because
+        # we added a subdirectory "test"
+        assert mpl.rcParams["savefig.directory"] == f"{tmp_path}/test"
diff --git a/src/_macosx.m b/src/_macosx.m
@@ -1022,33 +1022,24 @@ -(void)save_figure:(id)sender { gil_call_method(toolbar, "save_figure"); }
 {
     int result;
     const char* title;
+    const char* directory;
     const char* default_filename;
-    if (!PyArg_ParseTuple(args, "ss", &title, &default_filename)) {
+    if (!PyArg_ParseTuple(args, "sss", &title, &directory, &default_filename)) {
         return NULL;
     }
     NSSavePanel* panel = [NSSavePanel savePanel];
-    [panel setTitle: [NSString stringWithCString: title
-                                        encoding: NSASCIIStringEncoding]];
-    NSString* ns_default_filename =
-        [[NSString alloc]
-         initWithCString: default_filename
-         encoding: NSUTF8StringEncoding];
-    [panel setNameFieldStringValue: ns_default_filename];
+    [panel setTitle: [NSString stringWithUTF8String: title]];
+    [panel setDirectoryURL: [NSURL fileURLWithPath: [NSString stringWithUTF8String: directory]
+                                       isDirectory: YES]];
+    [panel setNameFieldStringValue: [NSString stringWithUTF8String: default_filename]];
     result = [panel runModal];
-    [ns_default_filename release];
     if (result == NSModalResponseOK) {
-        NSURL* url = [panel URL];
-        NSString* filename = [url path];
+        NSString *filename = [[panel URL] path];
         if (!filename) {
             PyErr_SetString(PyExc_RuntimeError, "Failed to obtain filename");
             return 0;
         }
-        unsigned int n = [filename length];
-        unichar* buffer = malloc(n*sizeof(unichar));
-        [filename getCharacters: buffer];
-        PyObject* string = PyUnicode_FromKindAndData(PyUnicode_2BYTE_KIND, buffer, n);
-        free(buffer);
-        return string;
+        return PyUnicode_FromString([filename UTF8String]);
     }
     Py_RETURN_NONE;
 }
