diff --git a/requirements.txt b/requirements.txt index 2666c38..2740e3b 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,3 +1,4 @@ +Pillow lxml>=4.1.1 requests>=2.20.0 diff --git a/safaribooks.py b/safaribooks.py index d66730b..022b053 100644 --- a/safaribooks.py +++ b/safaribooks.py @@ -13,6 +13,7 @@ from random import random from multiprocessing import Process, Queue, Value from urllib.parse import urljoin, urlsplit, urlparse +from PIL import Image PATH = os.path.dirname(os.path.realpath(__file__)) @@ -819,14 +820,35 @@ def _thread_download_images(self, url): stream=True) if response == 0: self.display.error("Error trying to retrieve this image: %s\n From: %s" % (image_name, url)) - - with open(image_path, 'wb') as img: + + with open(image_path, 'wb') as image_file: for chunk in response.iter_content(1024): - img.write(chunk) - + image_file.write(chunk) + try: + self._resize_image(image_path) + except: + # There are some image file which cannot open. Should we delete it? + pass + self.images_done_queue.put(1) self.display.state(len(self.images), self.images_done_queue.qsize()) - + + def _resize_image(self, image_path): + image_max_size = self.args.image_max_size + image_quality = self.args.image_quality + if image_max_size == 0 or image_quality == 0: + # No resize or changing quality + return + image = Image.open(image_path) + + if image_max_size > 0: + # Resize image if it's bigger than (image_max_size x image_max_size) + image.thumbnail((image_max_size, image_max_size)) + if image_quality > 0: + image.save(image_path, quality=image_quality) + else: + image.save(image_path) + def _start_multiprocessing(self, operation, full_queue): if len(full_queue) > 5: for i in range(0, len(full_queue), 5): @@ -1007,6 +1029,14 @@ def create_epub(self): "--no-cookies", dest="no_cookies", action='store_true', help="Prevent your session data to be saved into `cookies.json` file." ) + arguments.add_argument( + "--image-max-size", metavar="", dest="image_max_size", default=False, + help="Resize image if the image width/height is large than max size. 0 (default) for no resize" + ) + arguments.add_argument( + "--image-quality", metavar="", dest="image_quality", default=False, + help="Compress JPEG images with given quality. 0 (default) for keeping the original image's quality" + ) arguments.add_argument( "--no-kindle", dest="no_kindle", action='store_true', help="Remove some CSS rules that block overflow on `table` and `pre` elements." @@ -1035,5 +1065,19 @@ def create_epub(self): else: if args_parsed.no_cookies: arguments.error("invalid option: `--no-cookies` is valid only if you use the `--cred` option") - + + if args_parsed.image_max_size: + if not args_parsed.image_max_size.isdigit(): + arguments.error("invalid size: %s" % args_parsed.image_max_size) + args_parsed.image_max_size = int(args_parsed.image_max_size) + else: + args_parsed.image_max_size = 0 + + if args_parsed.image_quality: + if not args_parsed.image_quality.isdigit(): + arguments.error("invalid quality: %s" % args_parsed.image_quality) + args_parsed.image_quality = int(args_parsed.image_quality) + else: + args_parsed.image_quality = 0 + SafariBooks(args_parsed)