From d32e7b583e750e8d760c9eb5cad69b94aab4298e Mon Sep 17 00:00:00 2001
From: Eddie Lebow <elebow@users.noreply.github.com>
Date: Fri, 11 Oct 2019 22:31:07 -0400
Subject: [PATCH] Add optional `ext` keyword argument to set_cover()

Previously, set_cover() would rely on the Python standard module
`imghdr` to determine the type of the specified cover image file.
Unfortunately, imghdr does not support SVG, so set_cover() could
not handle that type (which is permitted by the EPUB spec).

Detecting SVG files programmatically is nontrivial, and is generally
considered to require parsing the file source.

Instead, simply provide an option for the client to specify the file
type with the new keyword argument.
---
 mkepub/mkepub.py            |  6 ++++--
 mkepub/tests/test_mkepub.py | 31 +++++++++++++++++++++++++++++++
 2 files changed, 35 insertions(+), 2 deletions(-)

diff --git a/mkepub/mkepub.py b/mkepub/mkepub.py
index 14a74a0..ab1b341 100644
--- a/mkepub/mkepub.py
+++ b/mkepub/mkepub.py
@@ -106,9 +106,11 @@ def add_font(self, name, data):
         self.fonts.append(name)
         self._add_file(pathlib.Path('fonts') / name, data)
 
-    def set_cover(self, data):
+    def set_cover(self, data, ext=None):
         """Set the cover image to the given data."""
-        self._cover = 'cover.' + imghdr.what(None, h=data)
+        if not ext:
+            ext = imghdr.what(None, h=data)
+        self._cover = 'cover.' + ext
         self._add_file(pathlib.Path('covers') / self._cover, data)
         self._write('cover.xhtml', 'EPUB/cover.xhtml', cover=self._cover)
 
diff --git a/mkepub/tests/test_mkepub.py b/mkepub/tests/test_mkepub.py
index 6535251..ac0d43c 100644
--- a/mkepub/tests/test_mkepub.py
+++ b/mkepub/tests/test_mkepub.py
@@ -6,6 +6,7 @@
 import mkepub
 import pathlib
 import pytest
+from unittest import mock
 
 ###############################################################################
 
@@ -104,3 +105,33 @@ def test_add_file():
     with open('mkepub/tests/cover.jpg', 'rb') as file:
         book._add_file('files/cover_1.jpg', file.read())
     assert (book.path / 'EPUB/files/cover_1.jpg').exists()
+
+
+@mock.patch("mkepub.Book._add_file", mock.MagicMock())
+@mock.patch("mkepub.Book._write", mock.MagicMock())
+class TestSetCover:
+    def test_set_cover_auto(self):
+        book = mkepub.Book('SetCover')
+
+        with open('mkepub/tests/cover.jpg', 'rb') as file:
+            data = file.read()
+        book.set_cover(data)
+        mkepub.Book._add_file.assert_called_with(pathlib.PosixPath('covers/cover.jpeg'), data)
+        mkepub.Book._write.assert_called_with(mock.ANY, mock.ANY, cover='cover.jpeg')
+
+    def test_set_cover_override(self):
+        book = mkepub.Book('SetCover')
+
+        with open('mkepub/tests/cover.jpg', 'rb') as file:
+            data = file.read()
+        book.set_cover(data, ext='png')
+        mkepub.Book._add_file.assert_called_with(pathlib.PosixPath('covers/cover.png'), data)
+        mkepub.Book._write.assert_called_with(mock.ANY, mock.ANY, cover='cover.png')
+
+    def test_set_cover_undetectable(self):
+        book = mkepub.Book('SetCover')
+
+        data = b'<svg></svg>'
+        book.set_cover(data, ext='svg')
+        mkepub.Book._add_file.assert_called_with(pathlib.PosixPath('covers/cover.svg'), data)
+        mkepub.Book._write.assert_called_with(mock.ANY, mock.ANY, cover='cover.svg')