indico_toolkit 7.2.1__tar.gz → 7.2.2__tar.gz

{indico_toolkit-7.2.1 → indico_toolkit-7.2.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: indico-toolkit
-Version: 7.2.1
+Version: 7.2.2
 Summary: Classes, functions, and abstractions for Indico IPA
 License: MIT
 Author: Indico Solutions Engineering

{indico_toolkit-7.2.1 → indico_toolkit-7.2.2}/indico_toolkit/__init__.py

@@ -21,4 +21,4 @@ __all__ = (
     "ToolkitStaggeredLoopError",
     "ToolkitStatusError",
 )
-__version__ = "7.2.1"
+__version__ = "7.2.2"

{indico_toolkit-7.2.1 → indico_toolkit-7.2.2}/indico_toolkit/etloutput/__init__.py

@@ -1,13 +1,13 @@
 from typing import TYPE_CHECKING, TypeAlias, TypeVar
 
-from ..results import NULL_BOX, NULL_SPAN, Box, Span
-from ..results.utils import get, has, json_loaded, str_decoded
-from .cell import Cell, CellType
-from .errors import EtlOutputError, TableCellNotFoundError, TokenNotFoundError
+from .box import NULL_BOX, Box
+from .cell import NULL_CELL, Cell, CellType
 from .etloutput import EtlOutput
-from .range import Range
-from .table import Table
-from .token import Token
+from .range import NULL_RANGE, Range
+from .span import NULL_SPAN, Span
+from .table import NULL_TABLE, Table
+from .token import NULL_TOKEN, Token
+from .utils import get, has, json_loaded, str_decoded
 
 if TYPE_CHECKING:
     from collections.abc import Awaitable, Callable
@@ -17,17 +17,18 @@ __all__ = (
     "Cell",
     "CellType",
     "EtlOutput",
-    "EtlOutputError",
     "load",
     "load_async",
     "NULL_BOX",
+    "NULL_CELL",
+    "NULL_RANGE",
     "NULL_SPAN",
+    "NULL_TABLE",
+    "NULL_TOKEN",
     "Range",
     "Span",
     "Table",
-    "TableCellNotFoundError",
     "Token",
-    "TokenNotFoundError",
 )
 
 Loadable: TypeAlias = "dict[str, object] | list[object] | str | bytes"

{indico_toolkit-7.2.1/indico_toolkit/results/predictions → indico_toolkit-7.2.2/indico_toolkit/etloutput}/box.py

@@ -1,10 +1,7 @@
 from dataclasses import dataclass
-from typing import TYPE_CHECKING
+from typing import Final
 
-from ..utils import get
-
-if TYPE_CHECKING:
-    from typing import Final
+from .utils import get
 
 
 @dataclass(frozen=True)
@@ -18,6 +15,30 @@ class Box:
     def __bool__(self) -> bool:
         return self != NULL_BOX
 
+    def __and__(self, other: "Box") -> "Box":
+        """
+        Return a new `Box` for the overlap between `self` and `other`
+        or `NULL_BOX` if they don't overlap.
+
+        Supports set-like `extraction.box & cell.box` syntax.
+        """
+        if (
+            self.page != other.page
+            or self.bottom <= other.top  # `self` is above `other`
+            or self.top >= other.bottom  # `self` is below `other`
+            or self.right <= other.left  # `self` is to the left of `other`
+            or self.left >= other.right  # `self` is to the right of `other`
+        ):
+            return NULL_BOX
+        else:
+            return Box(
+                page=self.page,
+                top=max(self.top, other.top),
+                left=max(self.left, other.left),
+                right=min(self.right, other.right),
+                bottom=min(self.bottom, other.bottom),
+            )
+
     def __lt__(self, other: "Box") -> bool:
         """
         Bounding boxes are sorted with vertical hysteresis. Those on the same line are
@@ -58,4 +79,4 @@ class Box:
 # object rather than using `None` or raising an error. This lets you e.g. sort by the
 # `box` attribute without having to constantly check for `None`, while still allowing
 # you do a "None check" with `bool(extraction.box)` or `extraction.box == NULL_BOX`.
-NULL_BOX: "Final" = Box(page=0, top=0, left=0, right=0, bottom=0)
+NULL_BOX: Final = Box(page=0, top=0, left=0, right=0, bottom=0)

{indico_toolkit-7.2.1 → indico_toolkit-7.2.2}/indico_toolkit/etloutput/cell.py

@@ -1,9 +1,11 @@
 from dataclasses import dataclass
 from enum import Enum
+from typing import Final
 
-from ..results import NULL_SPAN, Box, Span
-from ..results.utils import get
-from .range import Range
+from .box import NULL_BOX, Box
+from .range import NULL_RANGE, Range
+from .span import NULL_SPAN, Span
+from .utils import get
 
 
 class CellType(Enum):
@@ -19,6 +21,18 @@ class Cell:
     range: Range
     spans: "tuple[Span, ...]"
 
+    def __bool__(self) -> bool:
+        return self != NULL_CELL
+
+    def __hash__(self) -> int:
+        """
+        Uniquely identify cells by hashing their bounding box and spans.
+
+        This is small speedup for `.groupby(attrgetter("cell"))` compared to
+        dataclasses's default __hash__ implementation.
+        """
+        return hash((self.box, self.spans))
+
     @property
     def span(self) -> Span:
         """
@@ -45,3 +59,16 @@ class Cell:
             range=Range.from_dict(cell),
             spans=tuple(map(Span.from_dict, get(cell, list, "doc_offsets"))),
         )
+
+
+# It's more ergonomic to represent the lack of cells with a special null cell object
+# rather than using `None` or raising an error. This lets you e.g. sort by the `cell`
+# attribute without having to constantly check for `None`, while still allowing you do
+# a "None check" with `bool(extraction.cell)` or `extraction.cell == NULL_CELL`.
+NULL_CELL: Final = Cell(
+    type=CellType.CONTENT,
+    text="",
+    box=NULL_BOX,
+    range=NULL_RANGE,
+    spans=tuple(),
+)

{indico_toolkit-7.2.1 → indico_toolkit-7.2.2}/indico_toolkit/etloutput/etloutput.py

@@ -1,18 +1,20 @@
 import itertools
 from bisect import bisect_left, bisect_right
+from collections import namedtuple
 from dataclasses import dataclass
+from functools import cached_property
 from operator import attrgetter
 from typing import TYPE_CHECKING
 
-from ..results import Box, Span
-from .errors import TableCellNotFoundError, TokenNotFoundError
+from .box import Box
 from .table import Table
-from .token import Token
+from .token import NULL_TOKEN, Token
 
 if TYPE_CHECKING:
-    from collections.abc import Iterable
+    from collections.abc import Iterable, Iterator
 
     from .cell import Cell
+    from .span import Span
 
 
 @dataclass(frozen=True)
@@ -54,18 +56,19 @@ class EtlOutput:
             tables_on_page=table_pages,
         )
 
-    def token_for(self, span: Span) -> Token:
+    def token_for(self, span: "Span") -> Token:
         """
-        Return a `Token` that contains every character from `span`.
-        Raise `TokenNotFoundError` if one can't be produced.
+        Return a `Token` that contains every character from `span`
+        or `NULL_TOKEN` if one doesn't exist.
         """
         try:
             tokens = self.tokens_on_page[span.page]
             first = bisect_right(tokens, span.start, key=attrgetter("span.end"))
             last = bisect_left(tokens, span.end, lo=first, key=attrgetter("span.start"))
             tokens = tokens[first:last]
-        except (IndexError, ValueError) as error:
-            raise TokenNotFoundError(f"no token contains {span!r}") from error
+            assert tokens
+        except (AssertionError, IndexError, ValueError):
+            return NULL_TOKEN
 
         return Token(
             text=self.text[span.slice],
@@ -79,28 +82,53 @@ class EtlOutput:
             span=span,
         )
 
-    def table_cell_for(self, token: Token) -> "tuple[Table, Cell]":
+    _TableCellSpan = namedtuple("_TableCellSpan", ["table", "cell", "span"])
+
+    @cached_property
+    def _table_cell_spans_on_page(self) -> "tuple[tuple[_TableCellSpan, ...], ...]":
+        """
+        Order table cells on each page by their spans such that they can be bisected.
+        """
+        return tuple(
+            tuple(
+                sorted(
+                    (
+                        self._TableCellSpan(table, cell, span)
+                        for table in page_tables
+                        for cell in table.cells
+                        for span in cell.spans
+                        if span
+                    ),
+                    key=attrgetter("span"),
+                )
+            )
+            for page_tables in self.tables_on_page
+        )
+
+    def table_cells_for(self, span: "Span") -> "Iterator[tuple[Table, Cell]]":
         """
-        Return the `Table` and `Cell` that contain the midpoint of `token`.
-        Raise `TableCellNotFoundError` if it's not inside a table cell.
+        Yield the table cells that overlap with `span`.
+
+        Note: a single span may overlap the same cell multiple times causing it to be
+        yielded multiple times. Deduplication in `DocumentExtraction.table_cells`
+        accounts for this when OCR is assigned with `PredictionList.assign_ocr()`.
         """
-        token_vmid = (token.box.top + token.box.bottom) // 2
-        token_hmid = (token.box.left + token.box.right) // 2
-
-        for table in self.tables_on_page[token.box.page]:
-            if (
-                (table.box.top  <= token_vmid <= table.box.bottom) and
-                (table.box.left <= token_hmid <= table.box.right)
-            ):  # fmt: skip
-                break
-        else:
-            raise TableCellNotFoundError(f"no table contains {token!r}")
-
-        for cell in table.cells:
-            if (
-                (cell.box.top  <= token_vmid <= cell.box.bottom) and
-                (cell.box.left <= token_hmid <= cell.box.right)
-            ):  # fmt: skip
-                return table, cell
-        else:
-            raise TableCellNotFoundError(f"no cell contains {token!r}")
+        try:
+            page_table_cell_spans = self._table_cell_spans_on_page[span.page]
+            first = bisect_right(
+                page_table_cell_spans,
+                span.start,
+                key=attrgetter("span.end"),
+            )
+            last = bisect_left(
+                page_table_cell_spans,
+                span.end,
+                lo=first,
+                key=attrgetter("span.start"),
+            )
+            table_cell_spans = page_table_cell_spans[first:last]
+        except (IndexError, ValueError):
+            table_cell_spans = tuple()
+
+        for table, cell, span in table_cell_spans:
+            yield table, cell

{indico_toolkit-7.2.1 → indico_toolkit-7.2.2}/indico_toolkit/etloutput/range.py

@@ -1,6 +1,7 @@
 from dataclasses import dataclass
+from typing import Final
 
-from ..results.utils import get
+from .utils import get
 
 
 @dataclass(order=True, frozen=True)
@@ -12,6 +13,9 @@ class Range:
     rows: "tuple[int, ...]"
     columns: "tuple[int, ...]"
 
+    def __bool__(self) -> bool:
+        return self != NULL_RANGE
+
     @staticmethod
     def from_dict(cell: object) -> "Range":
         """
@@ -28,3 +32,17 @@ class Range:
             rows=tuple(rows),
             columns=tuple(columns),
         )
+
+
+# It's more ergonomic to represent the lack of ranges with a special null range object
+# rather than using `None` or raising an error. This lets you e.g. sort by the `range`
+# attribute without having to constantly check for `None`, while still allowing you do
+# a "None check" with `bool(cell.range)` or `cell.range == NULL_RANGE`.
+NULL_RANGE: Final = Range(
+    row=0,
+    column=0,
+    rowspan=0,
+    columnspan=0,
+    rows=tuple(),
+    columns=tuple(),
+)

{indico_toolkit-7.2.1/indico_toolkit/results/predictions → indico_toolkit-7.2.2/indico_toolkit/etloutput}/span.py

@@ -1,10 +1,7 @@
 from dataclasses import dataclass
-from typing import TYPE_CHECKING
+from typing import Any, Final
 
-from ..utils import get
-
-if TYPE_CHECKING:
-    from typing import Any, Final
+from .utils import get
 
 
 @dataclass(order=True, frozen=True)
@@ -20,6 +17,26 @@ class Span:
     def __bool__(self) -> bool:
         return self != NULL_SPAN
 
+    def __and__(self, other: "Span") -> "Span":
+        """
+        Return a new `Span` for the overlap between `self` and `other`
+        or `NULL_SPAN` if they don't overlap.
+
+        Supports set-like `extraction.span & cell.span` syntax.
+        """
+        if (
+            self.page != other.page
+            or self.end <= other.start  # `self` is to the left of `other`
+            or self.start >= other.end  # `self` is to the right of `other`
+        ):
+            return NULL_SPAN
+        else:
+            return Span(
+                page=self.page,
+                start=max(self.start, other.start),
+                end=min(self.end, other.end),
+            )
+
     @staticmethod
     def from_dict(span: object) -> "Span":
         return Span(
@@ -40,4 +57,4 @@ class Span:
 # rather than using `None` or raising an error. This lets you e.g. sort by the `span`
 # attribute without having to constantly check for `None`, while still allowing you do
 # a "None check" with `bool(extraction.span)` or `extraction.span == NULL_SPAN`.
-NULL_SPAN: "Final" = Span(page=0, start=0, end=0)
+NULL_SPAN: Final = Span(page=0, start=0, end=0)

{indico_toolkit-7.2.1 → indico_toolkit-7.2.2}/indico_toolkit/etloutput/table.py

@@ -1,18 +1,40 @@
 from dataclasses import dataclass
 from operator import attrgetter
+from typing import Final
 
-from ..results import Box
-from ..results.utils import get
+from .box import NULL_BOX, Box
 from .cell import Cell
+from .span import NULL_SPAN, Span
+from .utils import get
 
 
 @dataclass(frozen=True)
 class Table:
     box: Box
+    spans: "tuple[Span, ...]"
     cells: "tuple[Cell, ...]"
     rows: "tuple[tuple[Cell, ...], ...]"
     columns: "tuple[tuple[Cell, ...], ...]"
 
+    def __bool__(self) -> bool:
+        return self != NULL_TABLE
+
+    def __hash__(self) -> int:
+        """
+        Uniquely identify tables by hashing their bounding box and spans.
+
+        This is an order of magnitude speedup for `.groupby(attrgetter("table"))`
+        compared to dataclasses's default __hash__ implementation.
+        """
+        return hash((self.box, self.spans))
+
+    @property
+    def span(self) -> Span:
+        """
+        Return the first `Span` the table covers or `NULL_SPAN` otherwise.
+        """
+        return self.spans[0] if self.spans else NULL_SPAN
+
     @staticmethod
     def from_dict(table: object) -> "Table":
         """
@@ -50,9 +72,26 @@ class Table:
             for column in range(column_count)
         )  # fmt: skip
 
+        for doc_offset in get(table, list, "doc_offsets"):
+            doc_offset["page_num"] = page
+
         return Table(
             box=Box.from_dict(get(table, dict, "position")),
+            spans=tuple(map(Span.from_dict, get(table, list, "doc_offsets"))),
             cells=cells,
             rows=rows,
             columns=columns,
         )
+
+
+# It's more ergonomic to represent the lack of tables with a special null table object
+# rather than using `None` or raising an error. This lets you e.g. group by the `table`
+# attribute without having to constantly check for `None`, while still allowing you do
+# a "None check" with `bool(extraction.table)` or `extraction.table == NULL_TABLE`.
+NULL_TABLE: Final = Table(
+    box=NULL_BOX,
+    spans=tuple(),
+    cells=tuple(),
+    rows=tuple(),
+    columns=tuple(),
+)

{indico_toolkit-7.2.1 → indico_toolkit-7.2.2}/indico_toolkit/etloutput/token.py

@@ -1,7 +1,9 @@
 from dataclasses import dataclass
+from typing import Final
 
-from ..results import Box, Span
-from ..results.utils import get
+from .box import NULL_BOX, Box
+from .span import NULL_SPAN, Span
+from .utils import get
 
 
 @dataclass(frozen=True)
@@ -10,6 +12,9 @@ class Token:
     box: Box
     span: Span
 
+    def __bool__(self) -> bool:
+        return self != NULL_TOKEN
+
     @staticmethod
     def from_dict(token: object) -> "Token":
         """
@@ -23,3 +28,14 @@ class Token:
             box=Box.from_dict(get(token, dict, "position")),
             span=Span.from_dict(get(token, dict, "doc_offset")),
         )
+
+
+# It's more ergonomic to represent the lack of tokens with a special null token object
+# rather than using `None` or raising an error. This lets you e.g. sort by the `token`
+# attribute without having to constantly check for `None`, while still allowing you do
+# a "None check" with `bool(extraction.token)` or `extraction.token == NULL_TOKEN`.
+NULL_TOKEN: Final = Token(
+    text="",
+    box=NULL_BOX,
+    span=NULL_SPAN,
+)

indico_toolkit-7.2.2/indico_toolkit/etloutput/utils.py

@@ -0,0 +1,69 @@
+import json
+from typing import Any, TypeVar
+
+Value = TypeVar("Value")
+
+
+def get(value: object, value_type: "type[Value]", *keys: "str | int") -> Value:
+    """
+    Return the value of type `value_type` obtained by traversing `value` using `keys`.
+    Raise an error if a key doesn't exist or the value has the wrong type.
+    """
+    for key in keys:
+        if isinstance(value, dict):
+            if key in value:
+                value = value[key]
+            else:
+                raise KeyError(f"{key!r} not in {value.keys()!r}")
+        elif isinstance(value, list):
+            if isinstance(key, int):
+                if 0 <= key < len(value):
+                    value = value[key]
+                else:
+                    raise IndexError(f"{key} out of range [0,{len(value)})")
+            else:
+                raise TypeError(f"list can't be indexed with {key!r}")
+        else:
+            raise TypeError(f"{type(value)} can't be traversed")
+
+    if isinstance(value, value_type):
+        return value
+    else:
+        raise TypeError(f"value `{value!r}` doesn't have type {value_type}")
+
+
+def has(value: object, value_type: "type[Value]", *keys: "str | int") -> bool:
+    """
+    Check if `value` can be traversed using `keys` to a value of type `value_type`.
+    """
+    for key in keys:
+        if isinstance(value, dict) and key in value:
+            value = value[key]
+        elif isinstance(value, list) and isinstance(key, int) and 0 <= key < len(value):  # fmt: skip  # noqa: E501
+            value = value[key]
+        else:
+            return False
+
+    return isinstance(value, value_type)
+
+
+def json_loaded(value: "Any") -> "Any":
+    """
+    Ensure `value` has been loaded as JSON.
+    """
+    value = str_decoded(value)
+
+    if isinstance(value, str):
+        value = json.loads(value)
+
+    return value
+
+
+def str_decoded(value: str | bytes) -> str:
+    """
+    Ensure `value` has been decoded to a string.
+    """
+    if isinstance(value, bytes):
+        value = value.decode()
+
+    return value

{indico_toolkit-7.2.1 → indico_toolkit-7.2.2}/indico_toolkit/results/__init__.py

@@ -1,13 +1,10 @@
 from typing import TYPE_CHECKING, TypeAlias, TypeVar, overload
 
+from ..etloutput import NULL_BOX, NULL_SPAN, Box, Span
 from .document import Document
-from .errors import ResultError
 from .predictionlist import PredictionList
 from .predictions import (
-    NULL_BOX,
     NULL_CITATION,
-    NULL_SPAN,
-    Box,
     Classification,
     DocumentExtraction,
     Extraction,
@@ -15,7 +12,6 @@ from .predictions import (
     FormExtractionType,
     Group,
     Prediction,
-    Span,
     Summarization,
     Unbundling,
 )
@@ -45,7 +41,6 @@ __all__ = (
     "Prediction",
     "PredictionList",
     "Result",
-    "ResultError",
     "Review",
     "ReviewType",
     "Span",

{indico_toolkit-7.2.1 → indico_toolkit-7.2.2}/indico_toolkit/results/normalization.py

@@ -1,12 +1,9 @@
 import re
-from typing import TYPE_CHECKING
+from typing import Any
 
 from .task import TaskType
 from .utils import get, has
 
-if TYPE_CHECKING:
-    from typing import Any
-
 
 def normalize_result_dict(result: "Any") -> None:
     """

{indico_toolkit-7.2.1 → indico_toolkit-7.2.2}/indico_toolkit/results/predictionlist.py

@@ -1,6 +1,7 @@
 from collections import defaultdict
+from itertools import chain
 from operator import attrgetter
-from typing import TYPE_CHECKING, List, TypeVar, overload
+from typing import TYPE_CHECKING, Any, Final, List, SupportsIndex, TypeVar, overload
 
 from .predictions import (
     Classification,
@@ -13,25 +14,24 @@ from .predictions import (
     Unbundling,
 )
 from .review import Review, ReviewType
-from .task import TaskType
 from .utils import nfilter
 
 if TYPE_CHECKING:
-    from collections.abc import Callable, Collection, Container, Iterable
-    from typing import Any, Final, SupportsIndex
+    from collections.abc import Callable, Collection, Container, Iterable, Mapping
 
     from typing_extensions import Self
 
+    from ..etloutput import EtlOutput
     from .document import Document
     from .result import Result
-    from .task import Task
+    from .task import Task, TaskType
 
 PredictionType = TypeVar("PredictionType", bound=Prediction)
 OfType = TypeVar("OfType", bound=Prediction)
 KeyType = TypeVar("KeyType")
 
 # Non-None sentinel value to support `PredictionList.where(review=None)`.
-REVIEW_UNSPECIFIED: "Final" = Review(
+REVIEW_UNSPECIFIED: Final = Review(
     id=None, reviewer_id=None, notes=None, rejected=None, type=None  # type: ignore[arg-type]
 )
 
@@ -84,6 +84,43 @@ class PredictionList(List[PredictionType]):
 
         return self
 
+    def assign_ocr(
+        self,
+        etl_outputs: "Mapping[Document, EtlOutput]",
+        *,
+        tokens: bool = True,
+        tables: bool = True,
+    ) -> "Self":
+        """
+        Assign OCR tokens, tables, and/or cells using `etl_outputs`.
+
+        Use `tokens` or `tables` to skip lookup and assignment of those attributes.
+        """
+        extractions_by_document = self.oftype(
+            DocumentExtraction,
+        ).groupby(
+            attrgetter("document"),
+        )
+
+        for document, extractions in extractions_by_document.items():
+            etl_output = etl_outputs[document]
+
+            for extraction in extractions:
+                if tokens:
+                    extraction.tokens = list(
+                        filter(
+                            None,
+                            map(etl_output.token_for, extraction.spans),
+                        )
+                    )
+
+                if tables:
+                    extraction.table_cells = chain.from_iterable(
+                        map(etl_output.table_cells_for, extraction.spans)
+                    )
+
+        return self
+
     def groupby(
         self, key: "Callable[[PredictionType], KeyType]"
     ) -> "dict[KeyType, Self]":