1"""Recognize image file formats based on their first few bytes."""
2
3from os import PathLike
4import warnings
5
6__all__ = ["what"]
7
8
9warnings._deprecated(__name__, remove=(3, 13))
10
11
12#-------------------------#
13# Recognize image headers #
14#-------------------------#
15
16def what(file, h=None):
17    f = None
18    try:
19        if h is None:
20            if isinstance(file, (str, PathLike)):
21                f = open(file, 'rb')
22                h = f.read(32)
23            else:
24                location = file.tell()
25                h = file.read(32)
26                file.seek(location)
27        for tf in tests:
28            res = tf(h, f)
29            if res:
30                return res
31    finally:
32        if f: f.close()
33    return None
34
35
36#---------------------------------#
37# Subroutines per image file type #
38#---------------------------------#
39
40tests = []
41
42def test_jpeg(h, f):
43    """JPEG data with JFIF or Exif markers; and raw JPEG"""
44    if h[6:10] in (b'JFIF', b'Exif'):
45        return 'jpeg'
46    elif h[:4] == b'\xff\xd8\xff\xdb':
47        return 'jpeg'
48
49tests.append(test_jpeg)
50
51def test_png(h, f):
52    if h.startswith(b'\211PNG\r\n\032\n'):
53        return 'png'
54
55tests.append(test_png)
56
57def test_gif(h, f):
58    """GIF ('87 and '89 variants)"""
59    if h[:6] in (b'GIF87a', b'GIF89a'):
60        return 'gif'
61
62tests.append(test_gif)
63
64def test_tiff(h, f):
65    """TIFF (can be in Motorola or Intel byte order)"""
66    if h[:2] in (b'MM', b'II'):
67        return 'tiff'
68
69tests.append(test_tiff)
70
71def test_rgb(h, f):
72    """SGI image library"""
73    if h.startswith(b'\001\332'):
74        return 'rgb'
75
76tests.append(test_rgb)
77
78def test_pbm(h, f):
79    """PBM (portable bitmap)"""
80    if len(h) >= 3 and \
81        h[0] == ord(b'P') and h[1] in b'14' and h[2] in b' \t\n\r':
82        return 'pbm'
83
84tests.append(test_pbm)
85
86def test_pgm(h, f):
87    """PGM (portable graymap)"""
88    if len(h) >= 3 and \
89        h[0] == ord(b'P') and h[1] in b'25' and h[2] in b' \t\n\r':
90        return 'pgm'
91
92tests.append(test_pgm)
93
94def test_ppm(h, f):
95    """PPM (portable pixmap)"""
96    if len(h) >= 3 and \
97        h[0] == ord(b'P') and h[1] in b'36' and h[2] in b' \t\n\r':
98        return 'ppm'
99
100tests.append(test_ppm)
101
102def test_rast(h, f):
103    """Sun raster file"""
104    if h.startswith(b'\x59\xA6\x6A\x95'):
105        return 'rast'
106
107tests.append(test_rast)
108
109def test_xbm(h, f):
110    """X bitmap (X10 or X11)"""
111    if h.startswith(b'#define '):
112        return 'xbm'
113
114tests.append(test_xbm)
115
116def test_bmp(h, f):
117    if h.startswith(b'BM'):
118        return 'bmp'
119
120tests.append(test_bmp)
121
122def test_webp(h, f):
123    if h.startswith(b'RIFF') and h[8:12] == b'WEBP':
124        return 'webp'
125
126tests.append(test_webp)
127
128def test_exr(h, f):
129    if h.startswith(b'\x76\x2f\x31\x01'):
130        return 'exr'
131
132tests.append(test_exr)
133
134#--------------------#
135# Small test program #
136#--------------------#
137
138def test():
139    import sys
140    recursive = 0
141    if sys.argv[1:] and sys.argv[1] == '-r':
142        del sys.argv[1:2]
143        recursive = 1
144    try:
145        if sys.argv[1:]:
146            testall(sys.argv[1:], recursive, 1)
147        else:
148            testall(['.'], recursive, 1)
149    except KeyboardInterrupt:
150        sys.stderr.write('\n[Interrupted]\n')
151        sys.exit(1)
152
153def testall(list, recursive, toplevel):
154    import sys
155    import os
156    for filename in list:
157        if os.path.isdir(filename):
158            print(filename + '/:', end=' ')
159            if recursive or toplevel:
160                print('recursing down:')
161                import glob
162                names = glob.glob(os.path.join(glob.escape(filename), '*'))
163                testall(names, recursive, 0)
164            else:
165                print('*** directory (use -r) ***')
166        else:
167            print(filename + ':', end=' ')
168            sys.stdout.flush()
169            try:
170                print(what(filename))
171            except OSError:
172                print('*** not found ***')
173
174if __name__ == '__main__':
175    test()
176