|
9 | 9 | from nose.tools import assert_equal |
10 | 10 |
|
11 | 11 | import numpy as np |
12 | | -import tempfile |
13 | 12 | import os |
| 13 | +import shutil |
| 14 | +import tempfile |
| 15 | + |
14 | 16 |
|
15 | 17 | def test_uses_per_path(): |
16 | 18 | id = transforms.Affine2D() |
@@ -52,28 +54,35 @@ def check(master_transform, paths, all_transforms, |
52 | 54 | check(id, paths, tforms, offsets, [], []) |
53 | 55 | check(id, paths, tforms, offsets, facecolors[0:1], edgecolors) |
54 | 56 |
|
| 57 | + |
55 | 58 | @cleanup |
56 | 59 | def test_get_default_filename(): |
57 | | - test_dir = tempfile.mkdtemp() |
58 | | - plt.rcParams['savefig.directory'] = test_dir |
59 | | - fig = plt.figure() |
60 | | - canvas = FigureCanvasBase(fig) |
61 | | - filename = canvas.get_default_filename() |
62 | | - assert_equal(filename, 'image.png') |
| 60 | + try: |
| 61 | + test_dir = tempfile.mkdtemp() |
| 62 | + plt.rcParams['savefig.directory'] = test_dir |
| 63 | + fig = plt.figure() |
| 64 | + canvas = FigureCanvasBase(fig) |
| 65 | + filename = canvas.get_default_filename() |
| 66 | + assert_equal(filename, 'image.png') |
| 67 | + finally: |
| 68 | + shutil.rmtree(test_dir) |
| 69 | + |
63 | 70 |
|
64 | 71 | @cleanup |
65 | | -def test_get_default_filename_already_existing(): |
| 72 | +def test_get_default_filename_already_exists(): |
66 | 73 | # From #3068: Suggest non-existing default filename |
67 | | - test_dir = tempfile.mkdtemp() |
68 | | - plt.rcParams['savefig.directory'] = test_dir |
69 | | - fig = plt.figure() |
70 | | - canvas = FigureCanvasBase(fig) |
71 | | - |
72 | | - # create 'image.png' in figure's save dir |
73 | | - with open(os.path.join(test_dir, 'image.png'), 'a'): |
74 | | - filename = canvas.get_default_filename() |
75 | | - assert_equal(filename, 'image-1.png') |
| 74 | + try: |
| 75 | + test_dir = tempfile.mkdtemp() |
| 76 | + plt.rcParams['savefig.directory'] = test_dir |
| 77 | + fig = plt.figure() |
| 78 | + canvas = FigureCanvasBase(fig) |
76 | 79 |
|
| 80 | + # create 'image.png' in figure's save dir |
| 81 | + with open(os.path.join(test_dir, 'image.png'), 'a'): |
| 82 | + filename = canvas.get_default_filename() |
| 83 | + assert_equal(filename, 'image-1.png') |
| 84 | + finally: |
| 85 | + shutil.rmtree(test_dir) |
77 | 86 |
|
78 | 87 | if __name__ == "__main__": |
79 | 88 | import nose |
|
0 commit comments