1"""
2    sphinx.util.images
3    ~~~~~~~~~~~~~~~~~~
4
5    Image utility functions for Sphinx.
6
7    :copyright: Copyright 2007-2021 by the Sphinx team, see AUTHORS.
8    :license: BSD, see LICENSE for details.
9"""
10
11import base64
12import imghdr
13from collections import OrderedDict
14from os import path
15from typing import IO, NamedTuple, Optional, Tuple
16
17import imagesize
18
19try:
20    from PIL import Image
21except ImportError:
22    Image = None
23
24mime_suffixes = OrderedDict([
25    ('.gif', 'image/gif'),
26    ('.jpg', 'image/jpeg'),
27    ('.png', 'image/png'),
28    ('.pdf', 'application/pdf'),
29    ('.svg', 'image/svg+xml'),
30    ('.svgz', 'image/svg+xml'),
31    ('.ai', 'application/illustrator'),
32])
33
34DataURI = NamedTuple('DataURI', [('mimetype', str),
35                                 ('charset', str),
36                                 ('data', bytes)])
37
38
39def get_image_size(filename: str) -> Optional[Tuple[int, int]]:
40    try:
41        size = imagesize.get(filename)
42        if size[0] == -1:
43            size = None
44        elif isinstance(size[0], float) or isinstance(size[1], float):
45            size = (int(size[0]), int(size[1]))
46
47        if size is None and Image:  # fallback to Pillow
48            with Image.open(filename) as im:
49                size = im.size
50
51        return size
52    except Exception:
53        return None
54
55
56def guess_mimetype_for_stream(stream: IO, default: Optional[str] = None) -> Optional[str]:
57    imgtype = imghdr.what(stream)
58    if imgtype:
59        return 'image/' + imgtype
60    else:
61        return default
62
63
64def guess_mimetype(filename: str = '', default: Optional[str] = None) -> Optional[str]:
65    _, ext = path.splitext(filename.lower())
66    if ext in mime_suffixes:
67        return mime_suffixes[ext]
68    elif path.exists(filename):
69        with open(filename, 'rb') as f:
70            return guess_mimetype_for_stream(f, default=default)
71
72    return default
73
74
75def get_image_extension(mimetype: str) -> Optional[str]:
76    for ext, _mimetype in mime_suffixes.items():
77        if mimetype == _mimetype:
78            return ext
79
80    return None
81
82
83def parse_data_uri(uri: str) -> Optional[DataURI]:
84    if not uri.startswith('data:'):
85        return None
86
87    # data:[<MIME-type>][;charset=<encoding>][;base64],<data>
88    mimetype = 'text/plain'
89    charset = 'US-ASCII'
90
91    properties, data = uri[5:].split(',', 1)
92    for prop in properties.split(';'):
93        if prop == 'base64':
94            pass  # skip
95        elif prop.startswith('charset='):
96            charset = prop[8:]
97        elif prop:
98            mimetype = prop
99
100    image_data = base64.b64decode(data)
101    return DataURI(mimetype, charset, image_data)
102
103
104def test_svg(h: bytes, f: IO) -> Optional[str]:
105    """An additional imghdr library helper; test the header is SVG's or not."""
106    try:
107        if '<svg' in h.decode().lower():
108            return 'svg+xml'
109    except UnicodeDecodeError:
110        pass
111
112    return None
113
114
115# install test_svg() to imghdr
116# refs: https://docs.python.org/3.6/library/imghdr.html#imghdr.tests
117imghdr.tests.append(test_svg)
118