Files
old-books/_pipeline/tests/test_structure.py
T

81 lines
4.3 KiB
Python

"""Chapter heading detection, including damaged OCR headings, contents pages and books without chapters."""
import os
import sys
import unittest
sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), ".."))
import structure
def unit(text, page=1):
return {"text": text, "pdf_page": page}
BODY = ("A line of ordinary body text that goes on for a while. " * 12)
class HeadingTests(unittest.TestCase):
def numbers(self, units):
return [(h["number"], h["title"]) for h in structure.find_headings(units)]
def test_clean_headings_still_work(self):
units = [unit("CHAPTER ONE\nThe Start\n\n" + BODY), unit(BODY), unit("CHAPTER TWO\n\nThe Middle\n\n" + BODY), unit("CHAPTER 3\nEnd\n\n" + BODY)]
self.assertEqual(self.numbers(units), [(1, "The Start"), (2, "The Middle"), (3, "End")])
def test_ocr_damaged_numerals_are_taken_in_sequence(self):
units = [unit("CHAPTER I.\nBreath\n\n" + BODY), unit("CHAPTER II.\nTheory\n\n" + BODY),
unit("CHAPTER I I r.\nMore\n\n" + BODY), unit("CHAPTER IV.\nFour\n\n" + BODY),
unit("CHAPTER X V\nFifteen\n\n" + BODY)]
got = structure.find_headings(units)
self.assertEqual([h["number"] for h in got], [1, 2, 3, 4, 5]) # "X V" is not 15 here: 5 is next, so it was a gap in the numerals
self.assertEqual([h.get("guessed", False) for h in got], [False, False, True, False, True])
def test_contents_entry_is_replaced_by_the_real_heading(self):
units = [unit("CONTENTS\nCHAPTER I.\nBreath\nCHAPTER II.\nTheory\n"), unit("CHAPTER I.\nBreath\n\n" + BODY), unit("CHAPTER II.\nTheory\n\n" + BODY)]
got = structure.find_headings(units)
self.assertEqual([(h["number"], h["unit"]) for h in got], [(1, 1), (2, 2)])
self.assertTrue(got[0]["replaced_earlier"])
def test_sentences_starting_with_chapter_are_not_headings(self):
units = [unit("CHAPTER ONE\nStart\n\n" + BODY), unit("CHAPTER of the book is long and winding, said the guide.\n" + BODY)]
self.assertEqual(len(structure.find_headings(units)), 1)
def test_no_headings_gives_one_synthetic_chapter_from_the_first_body_page(self):
units = [unit("Title\nBy Someone"), unit("Short"), unit(BODY, 3), unit(BODY, 4)]
got = structure.find_headings(units)
self.assertEqual(len(got), 1)
self.assertTrue(got[0]["synthetic"])
self.assertEqual((got[0]["unit"], got[0]["number"]), (2, 1))
def test_ocr_damaged_chapter_word_is_recognised(self):
units = [unit("CHAPTER I.\nOne\n\n" + BODY), unit("OUAPTER II.\nTwo\n\n" + BODY),
unit("CUAPTER TX.\nNine?\n\n" + BODY), unit("CIEAPTER NEL,\nTwelve?\n\n" + BODY)]
got = structure.find_headings(units)
self.assertEqual([h["number"] for h in got], [1, 2, 3, 4]) # unreadable numerals count up in sequence
self.assertFalse(structure.is_chapter_word("CHAPTERS"))
self.assertFalse(structure.is_chapter_word("CAPTAIN"))
self.assertFalse(structure.is_chapter_word("SCIENCE"))
def test_long_title_without_blank_line_keeps_only_the_upper_case_lines(self):
body = "the Yogis not only bring about desired mental qualities and properties by will-power coupled with rhythmic breathing, but they also develop"
units = [unit("CHAPTER XVI.\nYOGI SPIRITUAL BREATHING.\n" + body + "\n" + body + "\n\n" + BODY)]
h = structure.find_headings(units)[0]
self.assertEqual(h["title"], "Yogi Spiritual Breathing.")
self.assertEqual(h["title_end"], 2)
def test_contents_in_chapter_line_form(self):
units = [unit("CONTENTS\nCHAPTER I.\nBreath Is Life........ 8\nCHAPTER II.\nThe Four Methods of Respiration ..... 27\n"), unit("CHAPTER I.\nX\n\n" + BODY)]
toc = structure.toc_chapter_lines(units, 1)
self.assertEqual(toc[2][0], "The Four Methods of Respiration")
self.assertEqual(toc[1][0], "Breath Is Life")
self.assertEqual(structure.unknown_words("The Four Metiiods of Respiration"), 1)
self.assertEqual(structure.unknown_words("The Four Methods of Respiration"), 0)
def test_nothing_readable_gives_no_headings(self):
self.assertEqual(structure.find_headings([unit(""), unit("x")]), [])
if __name__ == "__main__":
unittest.main()