diff --git a/mobi/extract.py b/mobi/extract.py index 20c38be..1e403ac 100644 --- a/mobi/extract.py +++ b/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) diff --git a/tests/test_mobi.py b/tests/test_mobi.py index a3c1b88..5fc0351 100644 --- a/tests/test_mobi.py +++ b/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)