1from io import BytesIO
2
3import numpy as np
4
5from matplotlib.testing.decorators import image_comparison
6import matplotlib.pyplot as plt
7import matplotlib.path as mpath
8import matplotlib.patches as mpatches
9from matplotlib.ticker import FuncFormatter
10
11
12@image_comparison(['bbox_inches_tight'], remove_text=True,
13                  savefig_kwarg={'bbox_inches': 'tight'})
14def test_bbox_inches_tight():
15    #: Test that a figure saved using bbox_inches='tight' is clipped correctly
16    data = [[66386, 174296, 75131, 577908, 32015],
17            [58230, 381139, 78045, 99308, 160454],
18            [89135, 80552, 152558, 497981, 603535],
19            [78415, 81858, 150656, 193263, 69638],
20            [139361, 331509, 343164, 781380, 52269]]
21
22    col_labels = row_labels = [''] * 5
23
24    rows = len(data)
25    ind = np.arange(len(col_labels)) + 0.3  # the x locations for the groups
26    cell_text = []
27    width = 0.4     # the width of the bars
28    yoff = np.zeros(len(col_labels))
29    # the bottom values for stacked bar chart
30    fig, ax = plt.subplots(1, 1)
31    for row in range(rows):
32        ax.bar(ind, data[row], width, bottom=yoff, align='edge', color='b')
33        yoff = yoff + data[row]
34        cell_text.append([''])
35    plt.xticks([])
36    plt.xlim(0, 5)
37    plt.legend([''] * 5, loc=(1.2, 0.2))
38    fig.legend([''] * 5, bbox_to_anchor=(0, 0.2), loc='lower left')
39    # Add a table at the bottom of the axes
40    cell_text.reverse()
41    plt.table(cellText=cell_text, rowLabels=row_labels, colLabels=col_labels,
42              loc='bottom')
43
44
45@image_comparison(['bbox_inches_tight_suptile_legend'],
46                  remove_text=False, savefig_kwarg={'bbox_inches': 'tight'})
47def test_bbox_inches_tight_suptile_legend():
48    plt.plot(np.arange(10), label='a straight line')
49    plt.legend(bbox_to_anchor=(0.9, 1), loc='upper left')
50    plt.title('Axis title')
51    plt.suptitle('Figure title')
52
53    # put an extra long y tick on to see that the bbox is accounted for
54    def y_formatter(y, pos):
55        if int(y) == 4:
56            return 'The number 4'
57        else:
58            return str(y)
59    plt.gca().yaxis.set_major_formatter(FuncFormatter(y_formatter))
60
61    plt.xlabel('X axis')
62
63
64@image_comparison(['bbox_inches_tight_suptile_non_default.png'],
65                  remove_text=False, savefig_kwarg={'bbox_inches': 'tight'},
66                  tol=0.1)  # large tolerance because only testing clipping.
67def test_bbox_inches_tight_suptitle_non_default():
68    fig, ax = plt.subplots()
69    fig.suptitle('Booo', x=0.5, y=1.1)
70
71
72@image_comparison(['bbox_inches_tight_clipping'],
73                  remove_text=True, savefig_kwarg={'bbox_inches': 'tight'})
74def test_bbox_inches_tight_clipping():
75    # tests bbox clipping on scatter points, and path clipping on a patch
76    # to generate an appropriately tight bbox
77    plt.scatter(np.arange(10), np.arange(10))
78    ax = plt.gca()
79    ax.set_xlim([0, 5])
80    ax.set_ylim([0, 5])
81
82    # make a massive rectangle and clip it with a path
83    patch = mpatches.Rectangle([-50, -50], 100, 100,
84                               transform=ax.transData,
85                               facecolor='blue', alpha=0.5)
86
87    path = mpath.Path.unit_regular_star(5).deepcopy()
88    path.vertices *= 0.25
89    patch.set_clip_path(path, transform=ax.transAxes)
90    plt.gcf().artists.append(patch)
91
92
93@image_comparison(['bbox_inches_tight_raster'],
94                  remove_text=True, savefig_kwarg={'bbox_inches': 'tight'})
95def test_bbox_inches_tight_raster():
96    """Test rasterization with tight_layout"""
97    fig, ax = plt.subplots()
98    ax.plot([1.0, 2.0], rasterized=True)
99
100
101def test_only_on_non_finite_bbox():
102    fig, ax = plt.subplots()
103    ax.annotate("", xy=(0, float('nan')))
104    ax.set_axis_off()
105    # we only need to test that it does not error out on save
106    fig.savefig(BytesIO(), bbox_inches='tight', format='png')
107
108
109def test_tight_pcolorfast():
110    fig, ax = plt.subplots()
111    ax.pcolorfast(np.arange(4).reshape((2, 2)))
112    ax.set(ylim=(0, .1))
113    buf = BytesIO()
114    fig.savefig(buf, bbox_inches="tight")
115    buf.seek(0)
116    height, width, _ = plt.imread(buf).shape
117    # Previously, the bbox would include the area of the image clipped out by
118    # the axes, resulting in a very tall image given the y limits of (0, 0.1).
119    assert width > height
120
121
122def test_noop_tight_bbox():
123    from PIL import Image
124    x_size, y_size = (10, 7)
125    dpi = 100
126    # make the figure just the right size up front
127    fig = plt.figure(frameon=False, dpi=dpi, figsize=(x_size/dpi, y_size/dpi))
128    ax = plt.Axes(fig, [0., 0., 1., 1.])
129    fig.add_axes(ax)
130    ax.set_axis_off()
131    ax.xaxis.set_visible(False)
132    ax.yaxis.set_visible(False)
133
134    data = np.arange(x_size * y_size).reshape(y_size, x_size)
135    ax.imshow(data, rasterized=True)
136
137    # When a rasterized Artist is included, a mixed-mode renderer does
138    # additional bbox adjustment. It should also be a no-op, and not affect the
139    # next save.
140    fig.savefig(BytesIO(), bbox_inches='tight', pad_inches=0, format='pdf')
141
142    out = BytesIO()
143    fig.savefig(out, bbox_inches='tight', pad_inches=0)
144    out.seek(0)
145    im = np.asarray(Image.open(out))
146    assert (im[:, :, 3] == 255).all()
147    assert not (im[:, :, :3] == 255).all()
148    assert im.shape == (7, 10, 4)
149