topological sorting of definitions by caller hierarchy

image_prediction/image_extractor/extractors/parsable.py

@@ -12,7 +12,67 @@ from image_prediction.image_extractor.extractor import ImageExtractor, ImageMeta
 from image_prediction.info import Info
 from image_prediction.stitching.stitching import stitch_pairs
 
-rounder = rcompose(round, int)
+
+class ParsablePDFImageExtractor(ImageExtractor):
+    def __init__(self, verbose=False):
+        self.doc: fitz.fitz.Document = None
+        self.verbose = verbose
+
+    def extract(self, pdf: bytes, page_range: range = None):
+        self.doc = fitz.Document(stream=pdf)
+
+        pages = extract_pages(self.doc, page_range) if page_range else self.doc
+
+        image_metadata_pairs = chain.from_iterable(
+            map(
+                self.__process_images_on_page,
+                tqdm(pages, desc="Extracting", disable=not self.verbose, total=len(page_range) if page_range else None),
+            )
+        )
+
+        return image_metadata_pairs
+
+    def __process_images_on_page(self, page: fitz.fitz.Page):
+        images = get_images_on_page(self.doc, page)
+        metadata = get_metadata_for_images_on_page(page)
+        get_image_infos.cache_clear()
+
+        image_metadata_pairs = starmap(
+            ImageMetadataPair, filter(compose(all, curry(map)(truth)), zip(images, metadata))
+        )
+        image_metadata_pairs = stitch_pairs(list(image_metadata_pairs))
+
+        return image_metadata_pairs
+
+
+def extract_pages(doc, page_range):
+    page_range = range(page_range.start + 1, page_range.stop + 1)
+    pages = map(doc.load_page, page_range)
+    return pages
+
+
+def get_images_on_page(doc, page: fitz.Page):
+    image_infos = get_image_infos(page)
+    xrefs = map(itemgetter("xref"), image_infos)
+    images = map(partial(load_image_from_xref, doc), xrefs)
+    return images
+
+
+def get_metadata_for_images_on_page(page: fitz.Page):
+    image_infos = get_image_infos(page)
+    metadata = map(get_image_metadata, image_infos)
+    metadata = map(partial(merge, get_page_metadata(page)), metadata)
+    return metadata
+
+
+def load_image_from_xref(doc, xref):
+    maybe_image = doc.extract_image(xref)
+    return Image.open(io.BytesIO(maybe_image["image"])) if maybe_image else None
+
+
+@lru_cache(maxsize=None)
+def get_image_infos(page: fitz.Page):
+    return page.get_image_info(xrefs=True)
 
 
 def get_image_metadata(image_info):
@@ -38,63 +98,4 @@ def get_page_metadata(page):
     }
 
 
-def extract_pages(doc, page_range):
+rounder = rcompose(round, int)
-    page_range = range(page_range.start + 1, page_range.stop + 1)
-    pages = map(doc.load_page, page_range)
-    return pages
-
-
-def load_image_from_xref(doc, xref):
-    maybe_image = doc.extract_image(xref)
-    return Image.open(io.BytesIO(maybe_image["image"])) if maybe_image else None
-
-
-@lru_cache(maxsize=None)
-def get_image_infos(page: fitz.Page):
-    return page.get_image_info(xrefs=True)
-
-
-def get_images_on_page(doc, page: fitz.Page):
-    image_infos = get_image_infos(page)
-    xrefs = map(itemgetter("xref"), image_infos)
-    images = map(partial(load_image_from_xref, doc), xrefs)
-    return images
-
-
-def get_metadata_for_images_on_page(page: fitz.Page):
-    image_infos = get_image_infos(page)
-    metadata = map(get_image_metadata, image_infos)
-    metadata = map(partial(merge, get_page_metadata(page)), metadata)
-    return metadata
-
-
-class ParsablePDFImageExtractor(ImageExtractor):
-    def __init__(self, verbose=False):
-        self.doc: fitz.fitz.Document = None
-        self.verbose = verbose
-
-    def __process_images_on_page(self, page: fitz.fitz.Page):
-        images = get_images_on_page(self.doc, page)
-        metadata = get_metadata_for_images_on_page(page)
-        get_image_infos.cache_clear()
-
-        image_metadata_pairs = starmap(
-            ImageMetadataPair, filter(compose(all, curry(map)(truth)), zip(images, metadata))
-        )
-        image_metadata_pairs = stitch_pairs(list(image_metadata_pairs))
-
-        return image_metadata_pairs
-
-    def extract(self, pdf: bytes, page_range: range = None):
-        self.doc = fitz.Document(stream=pdf)
-
-        pages = extract_pages(self.doc, page_range) if page_range else self.doc
-
-        image_metadata_pairs = chain.from_iterable(
-            map(
-                self.__process_images_on_page,
-                tqdm(pages, desc="Extracting", disable=not self.verbose, total=len(page_range) if page_range else None),
-            )
-        )
-
-        return image_metadata_pairs