Merge pull request #4221 from umairidris/default_filename

tacaswell · tacaswell · commit 51576d4930f1 · 2015-03-17T10:11:28.000-04:00
ENH : Suggest non-existing default filename in gui save window closes #3608
diff --git a/lib/matplotlib/backend_bases.py b/lib/matplotlib/backend_bases.py
@@ -2253,9 +2253,21 @@ def get_default_filename(self):
         Return a string, which includes extension, suitable for use as
         a default filename.
         """
-        default_filename = self.get_window_title() or 'image'
-        default_filename = default_filename.lower().replace(' ', '_')
-        return default_filename + '.' + self.get_default_filetype()
+        default_basename = self.get_window_title() or 'image'
+        default_basename = default_basename.lower().replace(' ', '_')
+        default_filetype = self.get_default_filetype()
+        default_filename = default_basename + '.' + default_filetype
+
+        save_dir = os.path.expanduser(rcParams.get('savefig.directory', ''))
+
+        # ensure non-existing filename in save dir
+        i = 1
+        while os.path.isfile(os.path.join(save_dir, default_filename)):
+            # attach numerical count to basename
+            default_filename = '{0}-{1}.{2}'.format(default_basename, i, default_filetype)
+            i += 1
+
+        return default_filename
 
     def switch_backends(self, FigureCanvasClass):
         """
diff --git a/lib/matplotlib/tests/test_backend_bases.py b/lib/matplotlib/tests/test_backend_bases.py
@@ -1,7 +1,17 @@
+from matplotlib.backend_bases import FigureCanvasBase
 from matplotlib.backend_bases import RendererBase
+from matplotlib.testing.decorators import image_comparison, cleanup
+
+import matplotlib.pyplot as plt
 import matplotlib.transforms as transforms
 import matplotlib.path as path
+
+from nose.tools import assert_equal
+
 import numpy as np
+import os
+import shutil
+import tempfile
 
 
 def test_uses_per_path():
@@ -43,3 +53,38 @@ def check(master_transform, paths, all_transforms,
     check(id, paths, tforms, offsets, facecolors, [])
     check(id, paths, tforms, offsets, [], [])
     check(id, paths, tforms, offsets, facecolors[0:1], edgecolors)
+
+
+@cleanup
+def test_get_default_filename():
+    try:
+        test_dir = tempfile.mkdtemp()
+        plt.rcParams['savefig.directory'] = test_dir
+        fig = plt.figure()
+        canvas = FigureCanvasBase(fig)
+        filename = canvas.get_default_filename()
+        assert_equal(filename, 'image.png')
+    finally:
+        shutil.rmtree(test_dir)
+
+
+@cleanup
+def test_get_default_filename_already_exists():
+    # From #3068: Suggest non-existing default filename
+    try:
+        test_dir = tempfile.mkdtemp()
+        plt.rcParams['savefig.directory'] = test_dir
+        fig = plt.figure()
+        canvas = FigureCanvasBase(fig)
+
+        # create 'image.png' in figure's save dir
+        open(os.path.join(test_dir, 'image.png'), 'w').close()
+
+        filename = canvas.get_default_filename()
+        assert_equal(filename, 'image-1.png')
+    finally:
+        shutil.rmtree(test_dir)
+
+if __name__ == "__main__":
+    import nose
+    nose.runmodule(argv=['-s', '--with-doctest'], exit=False)
