spec/fitzutils_spec.py

import os
import io

from mamba import description, it, before
from fitzutils import (
    open_pdf,
    ToCEntry,
    dump_toc
)
from pdftocio.tocparser import parse_toc

dirpath = os.path.dirname(os.path.abspath(__file__))

valid_file = os.path.join(dirpath, "files/level2.pdf")
invalid_file = os.path.join(dirpath, "files/nothing.pdf")

with description("open_pdf:") as self:
    with it("opens pdf file for reading"):
        with open_pdf(valid_file, False) as doc:
            assert doc is not None
            assert doc.page_count == 6

    with it("returns None if pdf file is invalid"):
        with open_pdf(invalid_file, False) as doc:
            assert doc is None

    with it("exits if pdf file is invalid and exit_on_error is true"):
        try:
            with open_pdf(invalid_file, True) as doc:
                assert False, "should have exited"
        except AssertionError as err:
            raise err
        except:
            pass

with description("ToCEntry") as self:
    with it("matches fitz's representation"):
        fitz_entry = [1, "title", 2]
        fitz_entry2 = [1, "title", 2, 100.0]

        toc_entry = ToCEntry(level=1, title="title", pagenum=2)
        toc_entry2 = ToCEntry(level=1, title="title", pagenum=2, vpos=100.0)

        assert toc_entry.to_fitz_entry() == fitz_entry
        assert toc_entry2.to_fitz_entry() == fitz_entry2

        assert ToCEntry(*fitz_entry) == toc_entry
        assert ToCEntry(*fitz_entry2) == toc_entry2

    with it("is sorted correctly"):
        entries = [
            ToCEntry(level=1, title="title4", pagenum=2, vpos=150.0),
            ToCEntry(level=1, title="title3", pagenum=2, vpos=90.0),
            ToCEntry(level=1, title="title5", pagenum=3, vpos=0.0),
            ToCEntry(level=1, title="title2", pagenum=1, vpos=150.0),
            ToCEntry(level=1, title="title1", pagenum=1, vpos=100.0),
            ToCEntry(level=1, title="title6", pagenum=5, vpos=200.0)
        ]

        expected = [
            ToCEntry(level=1, title="title1", pagenum=1, vpos=100.0),
            ToCEntry(level=1, title="title2", pagenum=1, vpos=150.0),
            ToCEntry(level=1, title="title3", pagenum=2, vpos=90.0),
            ToCEntry(level=1, title="title4", pagenum=2, vpos=150.0),
            ToCEntry(level=1, title="title5", pagenum=3, vpos=0.0),
            ToCEntry(level=1, title="title6", pagenum=5, vpos=200.0)
        ]
        assert sorted(entries, key=ToCEntry.key) == expected


with description("dump_toc") as self:
    with before.all:
        self.toc = [
            ToCEntry(level=1, title="title1", pagenum=1, vpos=100.0),
            ToCEntry(level=2, title="title2", pagenum=1, vpos=150.0),
            ToCEntry(level=3, title="title3", pagenum=2, vpos=90.0),
            ToCEntry(level=2, title="title4", pagenum=2, vpos=150.0),
            ToCEntry(level=2, title="title5", pagenum=3, vpos=0.0),
            ToCEntry(level=1, title="title6", pagenum=5, vpos=200.0)
        ]

        self.toc_novpos = [
            ToCEntry(level=1, title="title1", pagenum=1),
            ToCEntry(level=2, title="title2", pagenum=1),
            ToCEntry(level=3, title="title3", pagenum=2),
            ToCEntry(level=2, title="title4", pagenum=2),
            ToCEntry(level=2, title="title5", pagenum=3),
            ToCEntry(level=1, title="title6", pagenum=5)
        ]

    with it("won't print vpos if vpos is False"):
        toc_s = dump_toc(self.toc, False)
        f = io.StringIO(toc_s)
        assert parse_toc(f) == self.toc_novpos
        assert parse_toc(f) != self.toc

    with it("won't print vpos if vpos is missing"):
        toc_s = dump_toc(self.toc_novpos, True)
        f = io.StringIO(toc_s)
        assert parse_toc(f) == self.toc_novpos
        assert parse_toc(f) != self.toc