Merge remote-tracking branch 'image_formats/dev'

cps/helper.py

@@ -24,6 +24,7 @@ import ub
 from flask import current_app as app
 from tempfile import gettempdir
 import sys
+import io
 import os
 import re
 import unicodedata
@@ -34,6 +35,7 @@ from flask_babel import gettext as _
 from flask_login import current_user
 from babel.dates import format_datetime
 from datetime import datetime
+from PIL import Image
 import shutil
 import requests
 try:
@@ -442,27 +444,66 @@ def get_book_cover(cover_path):
         return send_from_directory(os.path.join(ub.config.config_calibre_dir, cover_path), "cover.jpg")
 
 
-# saves book cover to gdrive or locally
-def save_cover(url, book_path):
+# saves book cover from url
+def save_cover_from_url(url, book_path):
     img = requests.get(url)
-    if img.headers.get('content-type') != 'image/jpeg':
-        web.app.logger.error("Cover is no jpg file, can't save")
+    return save_cover(img, book_path)
+
+
+def save_cover_from_filestorage(filepath, saved_filename, img):
+    if hasattr(img,'_content'):
+        f = open(os.path.join(filepath, saved_filename), "wb")
+        f.write(img._content)
+        f.close()
+    else:
+        # check if file path exists, otherwise create it, copy file to calibre path and delete temp file
+        if not os.path.exists(filepath):
+            try:
+                os.makedirs(filepath)
+            except OSError:
+                web.app.logger.error(u"Failed to create path for cover")
+                return False
+        try:
+            img.save(os.path.join(filepath, saved_filename))
+        except OSError:
+            web.app.logger.error(u"Failed to store cover-file")
+            return False
+        except IOError:
+            web.app.logger.error(u"Cover-file is not a valid image file")
+            return False
+    return True
+
+
+# saves book cover to gdrive or locally
+def save_cover(img, book_path):
+    content_type = img.headers.get('content-type')
+    if content_type not in ('image/jpeg', 'image/png', 'image/webp'):
+        web.app.logger.error("Only jpg/jpeg/png/webp files are supported as coverfile")
         return False
 
+    # convert to jpg because calibre only supports jpg
+    if content_type in ('image/png', 'image/webp'):
+        if hasattr(img,'stream'):
+            imgc = Image.open(img.stream)
+        else:
+            imgc = Image.open(io.BytesIO(img.content))
+        im = imgc.convert('RGB')
+        tmp_bytesio = io.BytesIO()
+        im.save(tmp_bytesio, format='JPEG')
+        img._content = tmp_bytesio.getvalue()
+
     if ub.config.config_use_google_drive:
         tmpDir = gettempdir()
-        f = open(os.path.join(tmpDir, "uploaded_cover.jpg"), "wb")
-        f.write(img.content)
-        f.close()
-        gd.uploadFileToEbooksFolder(os.path.join(book_path, 'cover.jpg'), os.path.join(tmpDir, f.name))
-        web.app.logger.info("Cover is saved on Google Drive")
-        return True
+        if save_cover_from_filestorage(tmpDir, "uploaded_cover.jpg", img) is True:
+            gd.uploadFileToEbooksFolder(os.path.join(book_path, 'cover.jpg'),
+                                        os.path.join(tmpDir, "uploaded_cover.jpg"))
+            web.app.logger.info("Cover is saved on Google Drive")
+            return True
+        else:
+            return False
+    else:
+        return save_cover_from_filestorage(os.path.join(ub.config.config_calibre_dir, book_path), "cover.jpg", img)
 
-    f = open(os.path.join(ub.config.config_calibre_dir, book_path, "cover.jpg"), "wb")
-    f.write(img.content)
-    f.close()
-    web.app.logger.info("Cover is saved")
-    return True
 
 
 def do_download_file(book, book_format, data, headers):
@@ -487,7 +528,6 @@ def do_download_file(book, book_format, data, headers):
 
 
 
-
 def check_unrar(unrarLocation):
     error = False
     if os.path.exists(unrarLocation):