1from io import BytesIO
2from pathlib import Path
3
4import pytest
5
6from matplotlib.testing.decorators import image_comparison
7from matplotlib import pyplot as plt
8import matplotlib.cm as cm
9
10
11@image_comparison(['pngsuite.png'], tol=0.03)
12def test_pngsuite():
13    files = sorted(
14        (Path(__file__).parent / "baseline_images/pngsuite").glob("basn*.png"))
15
16    plt.figure(figsize=(len(files), 2))
17
18    for i, fname in enumerate(files):
19        data = plt.imread(fname)
20        cmap = None  # use default colormap
21        if data.ndim == 2:
22            # keep grayscale images gray
23            cmap = cm.gray
24        plt.imshow(data, extent=[i, i + 1, 0, 1], cmap=cmap)
25
26    plt.gca().patch.set_facecolor("#ddffff")
27    plt.gca().set_xlim(0, len(files))
28
29
30def test_truncated_file(tmpdir):
31    d = tmpdir.mkdir('test')
32    fname = str(d.join('test.png'))
33    fname_t = str(d.join('test_truncated.png'))
34    plt.savefig(fname)
35    with open(fname, 'rb') as fin:
36        buf = fin.read()
37    with open(fname_t, 'wb') as fout:
38        fout.write(buf[:20])
39
40    with pytest.raises(Exception):
41        plt.imread(fname_t)
42
43
44def test_truncated_buffer():
45    b = BytesIO()
46    plt.savefig(b)
47    b.seek(0)
48    b2 = BytesIO(b.read(20))
49    b2.seek(0)
50
51    with pytest.raises(Exception):
52        plt.imread(b2)
53