Limit image size, fixes #4464.
[librarian.git] / src / librarian / elements / figures / ilustr.py
1 # This file is part of Librarian, licensed under GNU Affero GPLv3 or later.
2 # Copyright © Fundacja Wolne Lektury. See NOTICE for more information.
3 #
4 import io
5 import urllib.parse
6 import urllib.request
7 from PIL import Image
8 from ..base import WLElement
9
10
11 MAX_PNG_WEIGHT = 200000
12
13
14 class Ilustr(WLElement):
15     SHOULD_HAVE_ID = True
16
17     EPUB_TAG = HTML_TAG = 'img'
18
19     def get_html_attr(self, builder):
20         ## TODO: thumbnail.
21
22         url = urllib.parse.urljoin(
23             builder.base_url,
24             self.get('src')
25         )
26         
27         imgfile = urllib.request.urlopen(url)
28         img = Image.open(imgfile)
29         th_format, ext, media_type = {
30             'GIF': ('GIF', 'gif', 'image/gif'),
31             'PNG': ('PNG', 'png', 'image/png'),
32         }.get(img.format, ('JPEG', 'jpg', 'image/jpeg'))
33
34         width = 600
35         if img.size[0] < width:
36             th = img
37         else:
38             th = img.resize((width, round(width * img.size[1] / img.size[0])))
39
40         buffer = io.BytesIO()
41         th.save(buffer, format=th_format)
42
43         # Limit PNG to 200K. If larger, convert to JPEG.
44         if th_format == 'PNG' and buffer.tell() > MAX_PNG_WEIGHT:
45             th_format, ext, media_type = 'JPEG', 'jpg', 'image/jpeg'
46             if th.mode != 'RGB':
47                 buffer = io.BytesIO()
48                 th = Image.alpha_composite(
49                     Image.new('RGBA', th.size, '#fff'),
50                     th.convert('RGBA')
51                 )
52                 th = th.convert('RGB')
53             th.save(buffer, format=th_format)
54
55         imgfile.close()
56         file_name = 'image%d.%s' % (
57             builder.assign_image_number(),
58             ext
59         )
60
61         builder.add_file(
62             content=buffer.getvalue(),
63             file_name=file_name,
64             media_type=media_type,
65         )
66         
67         return {
68             'src': file_name,
69             'alt': self.attrib.get('alt', ''),
70             'title': self.attrib.get('alt', ''),
71         }
72
73     get_epub_attr = get_html_attr