Initial commit of Roster
This commit is contained in:
Media
@@ -0,0 +1,511 @@
|
||||
import itertools
|
||||
import logging
|
||||
import re
|
||||
from collections import deque
|
||||
from dataclasses import asdict, dataclass, field
|
||||
from typing import (
|
||||
TYPE_CHECKING,
|
||||
Any,
|
||||
Callable,
|
||||
Dict,
|
||||
Iterable,
|
||||
Iterator,
|
||||
List,
|
||||
Optional,
|
||||
Pattern,
|
||||
Tuple,
|
||||
Union,
|
||||
)
|
||||
|
||||
from pdfminer.data_structures import NumberTree
|
||||
from pdfminer.pdfparser import PDFParser
|
||||
from pdfminer.pdftypes import PDFObjRef, resolve1
|
||||
from pdfminer.psparser import PSLiteral
|
||||
|
||||
from ._typing import T_bbox, T_obj
|
||||
from .utils import decode_text, geometry
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
if TYPE_CHECKING: # pragma: nocover
|
||||
from .page import Page
|
||||
from .pdf import PDF
|
||||
|
||||
|
||||
MatchFunc = Callable[["PDFStructElement"], bool]
|
||||
|
||||
|
||||
def _find_all(
|
||||
elements: Iterable["PDFStructElement"],
|
||||
matcher: Union[str, Pattern[str], MatchFunc],
|
||||
) -> Iterator["PDFStructElement"]:
|
||||
"""
|
||||
Common code for `find_all()` in trees and elements.
|
||||
"""
|
||||
|
||||
def match_tag(x: "PDFStructElement") -> bool:
|
||||
"""Match an element name."""
|
||||
return x.type == matcher
|
||||
|
||||
def match_regex(x: "PDFStructElement") -> bool:
|
||||
"""Match an element name by regular expression."""
|
||||
return matcher.match(x.type) # type: ignore
|
||||
|
||||
if isinstance(matcher, str):
|
||||
match_func = match_tag
|
||||
elif isinstance(matcher, re.Pattern):
|
||||
match_func = match_regex
|
||||
else:
|
||||
match_func = matcher # type: ignore
|
||||
d = deque(elements)
|
||||
while d:
|
||||
el = d.popleft()
|
||||
if match_func(el):
|
||||
yield el
|
||||
d.extendleft(reversed(el.children))
|
||||
|
||||
|
||||
class Findable:
|
||||
"""find() and find_all() methods that can be inherited to avoid
|
||||
repeating oneself"""
|
||||
|
||||
children: List["PDFStructElement"]
|
||||
|
||||
def find_all(
|
||||
self, matcher: Union[str, Pattern[str], MatchFunc]
|
||||
) -> Iterator["PDFStructElement"]:
|
||||
"""Iterate depth-first over matching elements in subtree.
|
||||
|
||||
The `matcher` argument is either an element name, a regular
|
||||
expression, or a function taking a `PDFStructElement` and
|
||||
returning `True` if the element matches.
|
||||
"""
|
||||
return _find_all(self.children, matcher)
|
||||
|
||||
def find(
|
||||
self, matcher: Union[str, Pattern[str], MatchFunc]
|
||||
) -> Optional["PDFStructElement"]:
|
||||
"""Find the first matching element in subtree.
|
||||
|
||||
The `matcher` argument is either an element name, a regular
|
||||
expression, or a function taking a `PDFStructElement` and
|
||||
returning `True` if the element matches.
|
||||
"""
|
||||
try:
|
||||
return next(_find_all(self.children, matcher))
|
||||
except StopIteration:
|
||||
return None
|
||||
|
||||
|
||||
@dataclass
|
||||
class PDFStructElement(Findable):
|
||||
type: str
|
||||
revision: Optional[int]
|
||||
id: Optional[str]
|
||||
lang: Optional[str]
|
||||
alt_text: Optional[str]
|
||||
actual_text: Optional[str]
|
||||
title: Optional[str]
|
||||
page_number: Optional[int]
|
||||
attributes: Dict[str, Any] = field(default_factory=dict)
|
||||
mcids: List[int] = field(default_factory=list)
|
||||
children: List["PDFStructElement"] = field(default_factory=list)
|
||||
|
||||
def __iter__(self) -> Iterator["PDFStructElement"]:
|
||||
return iter(self.children)
|
||||
|
||||
def all_mcids(self) -> Iterator[Tuple[Optional[int], int]]:
|
||||
"""Collect all MCIDs (with their page numbers, if there are
|
||||
multiple pages in the tree) inside a structure element.
|
||||
"""
|
||||
# Collect them depth-first to preserve ordering
|
||||
for mcid in self.mcids:
|
||||
yield self.page_number, mcid
|
||||
d = deque(self.children)
|
||||
while d:
|
||||
el = d.popleft()
|
||||
for mcid in el.mcids:
|
||||
yield el.page_number, mcid
|
||||
d.extendleft(reversed(el.children))
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
"""Return a compacted dict representation."""
|
||||
r = asdict(self)
|
||||
# Prune empty values (does not matter in which order)
|
||||
d = deque([r])
|
||||
while d:
|
||||
el = d.popleft()
|
||||
for k in list(el.keys()):
|
||||
if el[k] is None or el[k] == [] or el[k] == {}:
|
||||
del el[k]
|
||||
if "children" in el:
|
||||
d.extend(el["children"])
|
||||
return r
|
||||
|
||||
|
||||
class StructTreeMissing(ValueError):
|
||||
pass
|
||||
|
||||
|
||||
class PDFStructTree(Findable):
|
||||
"""Parse the structure tree of a PDF.
|
||||
|
||||
The constructor takes a `pdfplumber.PDF` and optionally a
|
||||
`pdfplumber.Page`. To avoid creating the entire tree for a large
|
||||
document it is recommended to provide a page.
|
||||
|
||||
This class creates a representation of the portion of the
|
||||
structure tree that reaches marked content sections, either for a
|
||||
single page, or for the whole document. Note that this is slightly
|
||||
different from the behaviour of other PDF libraries which will
|
||||
also include structure elements with no content.
|
||||
|
||||
If the PDF has no structure, the constructor will raise
|
||||
`StructTreeMissing`.
|
||||
|
||||
"""
|
||||
|
||||
page: Optional["Page"]
|
||||
|
||||
def __init__(self, doc: "PDF", page: Optional["Page"] = None):
|
||||
self.doc = doc.doc
|
||||
if "StructTreeRoot" not in self.doc.catalog:
|
||||
raise StructTreeMissing("PDF has no structure")
|
||||
self.root = resolve1(self.doc.catalog["StructTreeRoot"])
|
||||
self.role_map = resolve1(self.root.get("RoleMap", {}))
|
||||
self.class_map = resolve1(self.root.get("ClassMap", {}))
|
||||
self.children: List[PDFStructElement] = []
|
||||
|
||||
# If we have a specific page then we will work backwards from
|
||||
# its ParentTree - this is because structure elements could
|
||||
# span multiple pages, and the "Pg" attribute is *optional*,
|
||||
# so this is the approved way to get a page's structure...
|
||||
if page is not None:
|
||||
self.page = page
|
||||
self.pages = {page.page_number: page}
|
||||
self.page_dict = None
|
||||
# ...EXCEPT that the ParentTree is sometimes missing, in which
|
||||
# case we fall back to the non-approved way.
|
||||
parent_tree_obj = self.root.get("ParentTree")
|
||||
if parent_tree_obj is None:
|
||||
self._parse_struct_tree()
|
||||
else:
|
||||
parent_tree = NumberTree(parent_tree_obj)
|
||||
# If there is no marked content in the structure tree for
|
||||
# this page (which can happen even when there is a
|
||||
# structure tree) then there is no `StructParents`.
|
||||
# Note however that if there are XObjects in a page,
|
||||
# *they* may have `StructParent` (not `StructParents`)
|
||||
if "StructParents" not in self.page.page_obj.attrs:
|
||||
return
|
||||
parent_id = self.page.page_obj.attrs["StructParents"]
|
||||
# NumberTree should have a `get` method like it does in pdf.js...
|
||||
parent_array = resolve1(
|
||||
next(array for num, array in parent_tree.values if num == parent_id)
|
||||
)
|
||||
self._parse_parent_tree(parent_array)
|
||||
else:
|
||||
self.page = None
|
||||
# Overhead of creating pages shouldn't be too bad we hope!
|
||||
self.pages = {page.page_number: page for page in doc.pages}
|
||||
self.page_dict = {
|
||||
page.page_obj.pageid: page.page_number for page in self.pages.values()
|
||||
}
|
||||
self._parse_struct_tree()
|
||||
|
||||
def _make_attributes(
|
||||
self, obj: Dict[str, Any], revision: Optional[int]
|
||||
) -> Dict[str, Any]:
|
||||
attr_obj_list = []
|
||||
for key in "C", "A":
|
||||
if key not in obj:
|
||||
continue
|
||||
attr_obj = resolve1(obj[key])
|
||||
# It could be a list of attribute objects (why?)
|
||||
if isinstance(attr_obj, list):
|
||||
attr_obj_list.extend(attr_obj)
|
||||
else:
|
||||
attr_obj_list.append(attr_obj)
|
||||
attr_objs = []
|
||||
prev_obj = None
|
||||
for aref in attr_obj_list:
|
||||
# If we find a revision number, which might "follow the
|
||||
# revision object" (the spec is not clear about what this
|
||||
# should look like but it implies they are simply adjacent
|
||||
# in a flat array), then use it to decide whether to take
|
||||
# the previous object...
|
||||
if isinstance(aref, int):
|
||||
if aref == revision and prev_obj is not None:
|
||||
attr_objs.append(prev_obj)
|
||||
prev_obj = None
|
||||
else:
|
||||
if prev_obj is not None:
|
||||
attr_objs.append(prev_obj)
|
||||
prev_obj = resolve1(aref)
|
||||
if prev_obj is not None:
|
||||
attr_objs.append(prev_obj)
|
||||
# Now merge all the attribute objects in the collected to a
|
||||
# single set (again, the spec doesn't really explain this but
|
||||
# does say that attributes in /A supersede those in /C)
|
||||
attr = {}
|
||||
for obj in attr_objs:
|
||||
if isinstance(obj, PSLiteral):
|
||||
key = decode_text(obj.name)
|
||||
if key not in self.class_map:
|
||||
logger.warning("Unknown attribute class %s", key)
|
||||
continue
|
||||
obj = self.class_map[key]
|
||||
for k, v in obj.items():
|
||||
if isinstance(v, PSLiteral):
|
||||
attr[k] = decode_text(v.name)
|
||||
else:
|
||||
attr[k] = obj[k]
|
||||
return attr
|
||||
|
||||
def _make_element(self, obj: Any) -> Tuple[Optional[PDFStructElement], List[Any]]:
|
||||
# We hopefully caught these earlier
|
||||
assert "MCID" not in obj, "Uncaught MCR: %s" % obj
|
||||
assert "Obj" not in obj, "Uncaught OBJR: %s" % obj
|
||||
# Get page number if necessary
|
||||
page_number = None
|
||||
if self.page_dict is not None and "Pg" in obj:
|
||||
page_objid = obj["Pg"].objid
|
||||
assert page_objid in self.page_dict, "Object on unparsed page: %s" % obj
|
||||
page_number = self.page_dict[page_objid]
|
||||
obj_tag = ""
|
||||
if "S" in obj:
|
||||
obj_tag = decode_text(obj["S"].name)
|
||||
if obj_tag in self.role_map:
|
||||
obj_tag = decode_text(self.role_map[obj_tag].name)
|
||||
children = resolve1(obj["K"]) if "K" in obj else []
|
||||
if isinstance(children, int): # ugh... isinstance...
|
||||
children = [children]
|
||||
elif isinstance(children, dict): # a single object.. ugh...
|
||||
children = [obj["K"]]
|
||||
revision = obj.get("R")
|
||||
attributes = self._make_attributes(obj, revision)
|
||||
element_id = decode_text(resolve1(obj["ID"])) if "ID" in obj else None
|
||||
title = decode_text(resolve1(obj["T"])) if "T" in obj else None
|
||||
lang = decode_text(resolve1(obj["Lang"])) if "Lang" in obj else None
|
||||
alt_text = decode_text(resolve1(obj["Alt"])) if "Alt" in obj else None
|
||||
actual_text = (
|
||||
decode_text(resolve1(obj["ActualText"])) if "ActualText" in obj else None
|
||||
)
|
||||
element = PDFStructElement(
|
||||
type=obj_tag,
|
||||
id=element_id,
|
||||
page_number=page_number,
|
||||
revision=revision,
|
||||
lang=lang,
|
||||
title=title,
|
||||
alt_text=alt_text,
|
||||
actual_text=actual_text,
|
||||
attributes=attributes,
|
||||
)
|
||||
return element, children
|
||||
|
||||
def _parse_parent_tree(self, parent_array: List[Any]) -> None:
|
||||
"""Populate the structure tree using the leaves of the parent tree for
|
||||
a given page."""
|
||||
# First walk backwards from the leaves to the root, tracking references
|
||||
d = deque(parent_array)
|
||||
s = {}
|
||||
found_root = False
|
||||
while d:
|
||||
ref = d.popleft()
|
||||
# In the case where an MCID is not associated with any
|
||||
# structure, there will be a "null" in the parent tree.
|
||||
if ref == PDFParser.KEYWORD_NULL:
|
||||
continue
|
||||
if repr(ref) in s:
|
||||
continue
|
||||
obj = resolve1(ref)
|
||||
# This is required! It's in the spec!
|
||||
if "Type" in obj and decode_text(obj["Type"].name) == "StructTreeRoot":
|
||||
found_root = True
|
||||
else:
|
||||
# We hope that these are actual elements and not
|
||||
# references or marked-content sections...
|
||||
element, children = self._make_element(obj)
|
||||
# We have no page tree so we assume this page was parsed
|
||||
assert element is not None
|
||||
s[repr(ref)] = element, children
|
||||
d.append(obj["P"])
|
||||
# If we didn't reach the root something is quite wrong!
|
||||
assert found_root
|
||||
self._resolve_children(s)
|
||||
|
||||
def on_parsed_page(self, obj: Dict[str, Any]) -> bool:
|
||||
if "Pg" not in obj:
|
||||
return True
|
||||
page_objid = obj["Pg"].objid
|
||||
if self.page_dict is not None:
|
||||
return page_objid in self.page_dict
|
||||
if self.page is not None:
|
||||
# We have to do this to satisfy mypy
|
||||
if page_objid != self.page.page_obj.pageid:
|
||||
return False
|
||||
return True
|
||||
|
||||
def _parse_struct_tree(self) -> None:
|
||||
"""Populate the structure tree starting from the root, skipping
|
||||
unparsed pages and empty elements."""
|
||||
root = resolve1(self.root["K"])
|
||||
|
||||
# It could just be a single object ... it's in the spec (argh)
|
||||
if isinstance(root, dict):
|
||||
root = [self.root["K"]]
|
||||
d = deque(root)
|
||||
s = {}
|
||||
while d:
|
||||
ref = d.popleft()
|
||||
# In case the tree is actually a DAG and not a tree...
|
||||
if repr(ref) in s: # pragma: nocover (shouldn't happen)
|
||||
continue
|
||||
obj = resolve1(ref)
|
||||
# Deref top-level OBJR skipping refs to unparsed pages
|
||||
if isinstance(obj, dict) and "Obj" in obj:
|
||||
if not self.on_parsed_page(obj):
|
||||
continue
|
||||
ref = obj["Obj"]
|
||||
obj = resolve1(ref)
|
||||
element, children = self._make_element(obj)
|
||||
# Similar to above, delay resolving the children to avoid
|
||||
# tree-recursion.
|
||||
s[repr(ref)] = element, children
|
||||
for child in children:
|
||||
obj = resolve1(child)
|
||||
if isinstance(obj, dict):
|
||||
if not self.on_parsed_page(obj):
|
||||
continue
|
||||
if "Obj" in obj:
|
||||
child = obj["Obj"]
|
||||
elif "MCID" in obj:
|
||||
continue
|
||||
if isinstance(child, PDFObjRef):
|
||||
d.append(child)
|
||||
|
||||
# Traverse depth-first, removing empty elements (unsure how to
|
||||
# do this non-recursively)
|
||||
def prune(elements: List[Any]) -> List[Any]:
|
||||
next_elements = []
|
||||
for ref in elements:
|
||||
obj = resolve1(ref)
|
||||
if isinstance(ref, int):
|
||||
next_elements.append(ref)
|
||||
continue
|
||||
elif isinstance(obj, dict):
|
||||
if not self.on_parsed_page(obj):
|
||||
continue
|
||||
if "MCID" in obj:
|
||||
next_elements.append(obj["MCID"])
|
||||
continue
|
||||
elif "Obj" in obj:
|
||||
ref = obj["Obj"]
|
||||
element, children = s[repr(ref)]
|
||||
children = prune(children)
|
||||
# See assertions below
|
||||
if element is None or not children:
|
||||
del s[repr(ref)]
|
||||
else:
|
||||
s[repr(ref)] = element, children
|
||||
next_elements.append(ref)
|
||||
return next_elements
|
||||
|
||||
prune(root)
|
||||
self._resolve_children(s)
|
||||
|
||||
def _resolve_children(self, seen: Dict[str, Any]) -> None:
|
||||
"""Resolve children starting from the tree root based on references we
|
||||
saw when traversing the structure tree.
|
||||
"""
|
||||
root = resolve1(self.root["K"])
|
||||
# It could just be a single object ... it's in the spec (argh)
|
||||
if isinstance(root, dict):
|
||||
root = [self.root["K"]]
|
||||
self.children = []
|
||||
# Create top-level self.children
|
||||
parsed_root = []
|
||||
for ref in root:
|
||||
obj = resolve1(ref)
|
||||
if isinstance(obj, dict) and "Obj" in obj:
|
||||
if not self.on_parsed_page(obj):
|
||||
continue
|
||||
ref = obj["Obj"]
|
||||
if repr(ref) in seen:
|
||||
parsed_root.append(ref)
|
||||
d = deque(parsed_root)
|
||||
while d:
|
||||
ref = d.popleft()
|
||||
element, children = seen[repr(ref)]
|
||||
assert element is not None, "Unparsed element"
|
||||
for child in children:
|
||||
obj = resolve1(child)
|
||||
if isinstance(obj, int):
|
||||
element.mcids.append(obj)
|
||||
elif isinstance(obj, dict):
|
||||
# Skip out-of-page MCIDS and OBJRs
|
||||
if not self.on_parsed_page(obj):
|
||||
continue
|
||||
if "MCID" in obj:
|
||||
element.mcids.append(obj["MCID"])
|
||||
elif "Obj" in obj:
|
||||
child = obj["Obj"]
|
||||
# NOTE: if, not elif, in case of OBJR above
|
||||
if isinstance(child, PDFObjRef):
|
||||
child_element, _ = seen.get(repr(child), (None, None))
|
||||
if child_element is not None:
|
||||
element.children.append(child_element)
|
||||
d.append(child)
|
||||
self.children = [seen[repr(ref)][0] for ref in parsed_root]
|
||||
|
||||
def __iter__(self) -> Iterator[PDFStructElement]:
|
||||
return iter(self.children)
|
||||
|
||||
def element_bbox(self, el: PDFStructElement) -> T_bbox:
|
||||
"""Get the bounding box for an element for visual debugging."""
|
||||
page = None
|
||||
if self.page is not None:
|
||||
page = self.page
|
||||
elif el.page_number is not None:
|
||||
page = self.pages[el.page_number]
|
||||
bbox = el.attributes.get("BBox", None)
|
||||
if page is not None and bbox is not None:
|
||||
from .page import CroppedPage, _invert_box, _normalize_box
|
||||
|
||||
# Use secret knowledge of CroppedPage (cannot use
|
||||
# page.height because it is the *cropped* dimension, but
|
||||
# cropping does not actually translate coordinates)
|
||||
bbox = _invert_box(
|
||||
_normalize_box(bbox), page.mediabox[3] - page.mediabox[1]
|
||||
)
|
||||
# Use more secret knowledge of CroppedPage
|
||||
if isinstance(page, CroppedPage):
|
||||
rect = geometry.bbox_to_rect(bbox)
|
||||
rects = page._crop_fn([rect])
|
||||
if not rects:
|
||||
raise IndexError("Element no longer on page")
|
||||
return geometry.obj_to_bbox(rects[0])
|
||||
else:
|
||||
# Not sure why mypy complains here
|
||||
return bbox # type: ignore
|
||||
else:
|
||||
mcid_objs = []
|
||||
for page_number, mcid in el.all_mcids():
|
||||
objects: Iterable[T_obj]
|
||||
if page_number is None:
|
||||
if page is not None:
|
||||
objects = itertools.chain.from_iterable(page.objects.values())
|
||||
else:
|
||||
objects = [] # pragma: nocover
|
||||
else:
|
||||
objects = itertools.chain.from_iterable(
|
||||
self.pages[page_number].objects.values()
|
||||
)
|
||||
for c in objects:
|
||||
if c["mcid"] == mcid:
|
||||
mcid_objs.append(c)
|
||||
if not mcid_objs:
|
||||
raise IndexError("No objects found") # pragma: nocover
|
||||
return geometry.objects_to_bbox(mcid_objs)
|
||||
Reference in New Issue
Block a user