Allow selection of extraction type for extract

mobi/extract.py

@@ -1,14 +1,21 @@
 import shutil
 import tempfile
 from os.path import basename, exists, join, splitext
+from enum import IntFlag, auto
 
 from loguru import logger
 
 from mobi.kindleunpack import unpackBook
 
+class ExtractionType(IntFlag):
+    """The type of file format we will try to extract to. If ALL is chosen, the order will be: EPUB first, HTML second, PDF last."""
+    EPUB = auto()
+    HTML = auto()
+    PDF = auto()
+    ALL = EPUB | HTML | PDF
 
-def extract(infile):
-    """Extract mobi file and return path to epub file"""
+def extract(infile, extractable=ExtractionType.ALL):
+    """Extract mobi file and return path to ePub, HTML, or PDF file"""
 
     tempdir = tempfile.mkdtemp(prefix="mobiex")
     if hasattr(infile, "fileno"):
@@ -30,11 +37,11 @@ def extract(infile):
     epub_filepath = join(tempdir, "mobi8", fname_out_epub)
     html_filepath = join(tempdir, "mobi7", fname_out_html)
     pdf_filepath = join(tempdir, fname_out_pdf)
-    if exists(epub_filepath):
+    if exists(epub_filepath) and ExtractionType.EPUB in extractable:
         return tempdir, epub_filepath
-    elif exists(html_filepath):
+    elif exists(html_filepath) and ExtractionType.HTML in extractable:
         return tempdir, html_filepath
-    elif exists(pdf_filepath):
+    elif exists(pdf_filepath) and ExtractionType.PDF in extractable:
         return tempdir, pdf_filepath
     raise ValueError("Coud not extract from %s" % infile)
 

tests/test_mobi.py

@@ -3,6 +3,7 @@
 from os.path import abspath, dirname, exists, join, splitext
 
 import mobi
+from mobi.extract import ExtractionType
 
 TEST_DIR = dirname(abspath(__file__))
 
@@ -16,10 +17,24 @@ def test_extract():
             assert exists(filepath)
             shutil.rmtree(tempdir)
 
-
 def test_extract_file_like():
     with open(join(TEST_DIR, "demo.mobi"), "rb") as infile:
         tempdir, filepath = mobi.extract(infile)
         assert exists(tempdir)
         assert exists(filepath)
         shutil.rmtree(tempdir)
+
+def test_extract_specific_output():
+
+    extractiontypetofileextention = [
+        [ExtractionType.EPUB, ".epub"],
+        [ExtractionType.HTML, ".html"],
+    ]
+
+    with open(join(TEST_DIR, "demo.mobi"), "rb") as infile:
+        for ext in extractiontypetofileextention:
+            tempdir, filepath = mobi.extract(infile, extractable=ext[0])
+            assert exists(tempdir)
+            assert exists(filepath)
+            assert (ext[1] in filepath)
+            shutil.rmtree(tempdir)