indico_toolkit 6.14.0__py3-none-any.whl

indico_toolkit/etloutput/__init__.py

@@ -0,0 +1,210 @@
+from typing import TYPE_CHECKING
+
+from ..results import NULL_BOX, NULL_SPAN, Box, Span
+from ..results.utilities import get, has
+from .cell import Cell, CellType
+from .errors import EtlOutputError, TableCellNotFoundError, TokenNotFoundError
+from .etloutput import EtlOutput
+from .range import Range
+from .table import Table
+from .token import Token
+
+if TYPE_CHECKING:
+    from collections.abc import Awaitable, Callable
+    from typing import Any
+
+__all__ = (
+    "Box",
+    "Cell",
+    "CellType",
+    "EtlOutput",
+    "EtlOutputError",
+    "load",
+    "load_async",
+    "NULL_BOX",
+    "NULL_SPAN",
+    "Range",
+    "Span",
+    "Table",
+    "TableCellNotFoundError",
+    "Token",
+    "TokenNotFoundError",
+)
+
+
+def load(
+    etl_output_uri: str,
+    *,
+    reader: "Callable[..., Any]",
+    text: bool = True,
+    tokens: bool = True,
+    tables: bool = False,
+) -> EtlOutput:
+    """
+    Load `etl_output_uri` as an ETL Output dataclass. A `reader` function must be
+    supplied to read JSON files from disk, storage API, or Indico client.
+
+    Use `text`, `tokens`, and `tables` to specify what to load.
+
+    ```
+    result = results.load(submission.result_file, reader=read_uri)
+    etl_outputs = {
+        document: etloutput.load(document.etl_output_uri, reader=read_uri)
+        for document in result.documents
+        if not document.failed
+    }
+    ```
+    """
+    etl_output = reader(etl_output_uri)
+    tables_uri = etl_output_uri.replace("etl_output.json", "tables.json")
+
+    if has(etl_output, str, "pages", 0, "page_info"):
+        return _load_v1(etl_output, tables_uri, reader, text, tokens, tables)
+    else:
+        return _load_v3(etl_output, tables_uri, reader, text, tokens, tables)
+
+
+async def load_async(
+    etl_output_uri: str,
+    *,
+    reader: "Callable[..., Awaitable[Any]]",
+    text: bool = True,
+    tokens: bool = True,
+    tables: bool = False,
+) -> EtlOutput:
+    """
+    Load `etl_output_uri` as an ETL Output dataclass. A `reader` coroutine must be
+    supplied to read JSON files from disk, storage API, or Indico client.
+
+    Use `text`, `tokens`, and `tables` to specify what to load.
+
+    ```
+    result = await results.load_async(submission.result_file, reader=read_uri)
+    etl_outputs = {
+        document: await etloutput.load_async(document.etl_output_uri, reader=read_uri)
+        for document in result.documents
+        if not document.failed
+    }
+    ```
+    """
+    etl_output = await reader(etl_output_uri)
+    tables_uri = etl_output_uri.replace("etl_output.json", "tables.json")
+
+    if has(etl_output, str, "pages", 0, "page_info"):
+        return await _load_v1_async(
+            etl_output, tables_uri, reader, text, tokens, tables
+        )
+    else:
+        return await _load_v3_async(
+            etl_output, tables_uri, reader, text, tokens, tables
+        )
+
+
+def _load_v1(
+    etl_output: "Any",
+    tables_uri: str,
+    reader: "Callable[..., Any]",
+    text: bool,
+    tokens: bool,
+    tables: bool,
+) -> EtlOutput:
+    if text or tokens:
+        pages = tuple(
+            reader(get(page, str, "page_info"))
+            for page in get(etl_output, list, "pages")
+        )
+        text_by_page = map(lambda page: get(page, str, "pages", 0, "text"), pages)
+        tokens_by_page = map(lambda page: get(page, list, "tokens"), pages)
+    else:
+        text_by_page = ()  # type: ignore[assignment]
+        tokens_by_page = ()  # type: ignore[assignment]
+
+    if tables:
+        tables_by_page = reader(tables_uri)
+    else:
+        tables_by_page = ()
+
+    return EtlOutput.from_pages(text_by_page, tokens_by_page, tables_by_page)
+
+
+def _load_v3(
+    etl_output: "Any",
+    tables_uri: str,
+    reader: "Callable[..., Any]",
+    text: bool,
+    tokens: bool,
+    tables: bool,
+) -> EtlOutput:
+    pages = get(etl_output, list, "pages")
+
+    if text or tokens:
+        text_by_page = map(lambda page: reader(get(page, str, "text")), pages)
+    else:
+        text_by_page = ()  # type: ignore[assignment]
+
+    if tokens:
+        tokens_by_page = map(lambda page: reader(get(page, str, "tokens")), pages)
+    else:
+        tokens_by_page = ()  # type: ignore[assignment]
+
+    if tables:
+        tables_by_page = reader(tables_uri)
+    else:
+        tables_by_page = ()
+
+    return EtlOutput.from_pages(text_by_page, tokens_by_page, tables_by_page)
+
+
+async def _load_v1_async(
+    etl_output: "Any",
+    tables_uri: str,
+    reader: "Callable[..., Awaitable[Any]]",
+    text: bool,
+    tokens: bool,
+    tables: bool,
+) -> EtlOutput:
+    if text or tokens:
+        pages = [
+            await reader(get(page, str, "page_info"))
+            for page in get(etl_output, list, "pages")
+        ]
+        text_by_page = map(lambda page: get(page, str, "pages", 0, "text"), pages)
+        tokens_by_page = map(lambda page: get(page, list, "tokens"), pages)
+    else:
+        text_by_page = ()  # type: ignore[assignment]
+        tokens_by_page = ()  # type: ignore[assignment]
+
+    if tables:
+        tables_by_page = await reader(tables_uri)
+    else:
+        tables_by_page = ()
+
+    return EtlOutput.from_pages(text_by_page, tokens_by_page, tables_by_page)
+
+
+async def _load_v3_async(
+    etl_output: "Any",
+    tables_uri: str,
+    reader: "Callable[..., Awaitable[Any]]",
+    text: bool,
+    tokens: bool,
+    tables: bool,
+) -> EtlOutput:
+    pages = get(etl_output, list, "pages")
+
+    if text or tokens:
+        text_by_page = [await reader(get(page, str, "text")) for page in pages]
+    else:
+        text_by_page = ()  # type: ignore[assignment]
+
+    if tokens:
+        tokens_by_page = [await reader(get(page, str, "tokens")) for page in pages]
+    else:
+        tokens_by_page = ()  # type: ignore[assignment]
+
+    if tables:
+        tables_by_page = await reader(tables_uri)
+    else:
+        tables_by_page = ()
+
+    return EtlOutput.from_pages(text_by_page, tokens_by_page, tables_by_page)

indico_toolkit/etloutput/cell.py

@@ -0,0 +1,47 @@
+from dataclasses import dataclass
+from enum import Enum
+
+from ..results import NULL_SPAN, Box, Span
+from ..results.utilities import get
+from .range import Range
+
+
+class CellType(Enum):
+    HEADER = "header"
+    CONTENT = "content"
+
+
+@dataclass(frozen=True)
+class Cell:
+    type: CellType
+    text: str
+    box: Box
+    range: Range
+    spans: "tuple[Span, ...]"
+
+    @property
+    def span(self) -> Span:
+        """
+        Return the first `Span` the cell covers or `NULL_SPAN` otherwise.
+
+        Empty cells have no spans.
+        """
+        return self.spans[0] if self.spans else NULL_SPAN
+
+    @staticmethod
+    def from_dict(cell: object, page: int) -> "Cell":
+        """
+        Create a `Cell` from a v1 or v3 cell dictionary.
+        """
+        get(cell, dict, "position")["page_num"] = page
+
+        for doc_offset in get(cell, list, "doc_offsets"):
+            doc_offset["page_num"] = page
+
+        return Cell(
+            type=CellType(get(cell, str, "cell_type")),
+            text=get(cell, str, "text"),
+            box=Box.from_dict(get(cell, dict, "position")),
+            range=Range.from_dict(cell),
+            spans=tuple(map(Span.from_dict, get(cell, list, "doc_offsets"))),
+        )

indico_toolkit/etloutput/errors.py

@@ -0,0 +1,16 @@
+class EtlOutputError(Exception):
+    """
+    Raised when an error occurs while loading an ETL Output file.
+    """
+
+
+class TokenNotFoundError(EtlOutputError):
+    """
+    Raised when a Token can't be found for a Prediction.
+    """
+
+
+class TableCellNotFoundError(EtlOutputError):
+    """
+    Raised when a Table Cell can't be found for a Token.
+    """

indico_toolkit/etloutput/etloutput.py

@@ -0,0 +1,110 @@
+import itertools
+from bisect import bisect_left, bisect_right
+from dataclasses import dataclass
+from operator import attrgetter
+from typing import TYPE_CHECKING
+
+from ..results import Box, Span
+from .errors import TableCellNotFoundError, TokenNotFoundError
+from .table import Table
+from .token import Token
+
+if TYPE_CHECKING:
+    from collections.abc import Iterable
+
+    from .cell import Cell
+
+
+@dataclass(frozen=True)
+class EtlOutput:
+    text: str
+    text_on_page: "tuple[str, ...]"
+
+    tokens: "tuple[Token, ...]"
+    tokens_on_page: "tuple[tuple[Token, ...], ...]"
+
+    tables: "tuple[Table, ...]"
+    tables_on_page: "tuple[tuple[Table, ...], ...]"
+
+    @staticmethod
+    def from_pages(
+        text_by_page: "Iterable[str]",
+        token_dicts_by_page: "Iterable[Iterable[object]]",
+        table_dicts_by_page: "Iterable[Iterable[object]]",
+    ) -> "EtlOutput":
+        """
+        Create an `EtlOutput` from v1 or v3 page lists.
+        """
+        text_by_page = tuple(text_by_page)
+        tokens_by_page = tuple(
+            tuple(sorted(map(Token.from_dict, token_dict_page), key=attrgetter("span")))
+            for token_dict_page in token_dicts_by_page
+        )
+        tables_by_page = tuple(
+            tuple(sorted(map(Table.from_dict, table_dict_page), key=attrgetter("box")))
+            for table_dict_page in table_dicts_by_page
+        )
+
+        return EtlOutput(
+            text="\n".join(text_by_page),
+            text_on_page=text_by_page,
+            tokens=tuple(itertools.chain.from_iterable(tokens_by_page)),
+            tokens_on_page=tokens_by_page,
+            tables=tuple(itertools.chain.from_iterable(tables_by_page)),
+            tables_on_page=tables_by_page,
+        )
+
+    def token_for(self, span: Span) -> Token:
+        """
+        Return a `Token` that contains every character from `span`.
+        Raise `TokenNotFoundError` if one can't be produced.
+        """
+        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]
+
+            return Token(
+                text=self.text[span.slice],
+                span=span,
+                box=Box(
+                    page=min(token.box.page for token in tokens),
+                    top=min(token.box.top for token in tokens),
+                    left=min(token.box.left for token in tokens),
+                    right=max(token.box.right for token in tokens),
+                    bottom=max(token.box.bottom for token in tokens),
+                ),
+            )
+        except (IndexError, ValueError) as error:
+            raise TokenNotFoundError(f"no token contains {span!r}") from error
+
+    def table_cell_for(self, token: Token) -> "tuple[Table, Cell]":
+        """
+        Return the `Table` and `Cell` that contain the midpoint of `token`.
+        Raise `TableCellNotFoundError` if it's not inside a table cell.
+        """
+        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}")
+
+        try:
+            row_index = bisect_left(
+                table.rows, token_vmid, key=lambda row: row[0].box.bottom
+            )
+            row = table.rows[row_index]
+
+            cell_index = bisect_left(row, token_hmid, key=attrgetter("box.right"))
+            cell = row[cell_index]
+        except (IndexError, ValueError) as error:
+            raise TableCellNotFoundError(f"no cell contains {token!r}") from error
+
+        return table, cell

indico_toolkit/etloutput/range.py

@@ -0,0 +1,30 @@
+from dataclasses import dataclass
+
+from ..results.utilities import get
+
+
+@dataclass(order=True, frozen=True)
+class Range:
+    row: int
+    column: int
+    rowspan: int
+    columnspan: int
+    rows: "tuple[int, ...]"
+    columns: "tuple[int, ...]"
+
+    @staticmethod
+    def from_dict(cell: object) -> "Range":
+        """
+        Create a `Range` from a v1 or v3 cell dictionary.
+        """
+        rows = get(cell, list, "rows")
+        columns = get(cell, list, "columns")
+
+        return Range(
+            row=rows[0],
+            column=columns[0],
+            rowspan=len(rows),
+            columnspan=len(columns),
+            rows=tuple(rows),
+            columns=tuple(columns),
+        )

indico_toolkit/etloutput/table.py

@@ -0,0 +1,44 @@
+from dataclasses import dataclass
+from operator import attrgetter
+
+from ..results import Box
+from ..results.utilities import get
+from .cell import Cell
+
+
+@dataclass(frozen=True)
+class Table:
+    box: Box
+    cells: "tuple[Cell, ...]"
+    rows: "tuple[tuple[Cell, ...], ...]"
+    columns: "tuple[tuple[Cell, ...], ...]"
+
+    @staticmethod
+    def from_dict(table: object) -> "Table":
+        """
+        Create a `Table` from a v1 or v3 table dictionary.
+        """
+        page = get(table, int, "page_num")
+        get(table, dict, "position")["page_num"] = page
+
+        cells = tuple(
+            sorted(
+                (Cell.from_dict(cell, page) for cell in get(table, list, "cells")),
+                key=attrgetter("range"),
+            )
+        )
+        rows = tuple(
+            tuple(cell for cell in cells if row in cell.range.rows)
+            for row in range(get(table, int, "num_rows"))
+        )
+        columns = tuple(
+            tuple(cell for cell in cells if column in cell.range.columns)
+            for column in range(get(table, int, "num_columns"))
+        )
+
+        return Table(
+            box=Box.from_dict(get(table, dict, "position")),
+            cells=cells,
+            rows=rows,
+            columns=columns,
+        )

indico_toolkit/etloutput/token.py

@@ -0,0 +1,25 @@
+from dataclasses import dataclass
+
+from ..results import Box, Span
+from ..results.utilities import get
+
+
+@dataclass(frozen=True)
+class Token:
+    text: str
+    box: Box
+    span: Span
+
+    @staticmethod
+    def from_dict(token: object) -> "Token":
+        """
+        Create a `Token` from a v1 or v3 token dictionary.
+        """
+        get(token, dict, "position")["page_num"] = get(token, int, "page_num")
+        get(token, dict, "doc_offset")["page_num"] = get(token, int, "page_num")
+
+        return Token(
+            text=get(token, str, "text"),
+            box=Box.from_dict(get(token, dict, "position")),
+            span=Span.from_dict(get(token, dict, "doc_offset")),
+        )

indico_toolkit/indico_wrapper/__init__.py

@@ -0,0 +1,15 @@
+from .dataset import Datasets
+from .doc_extraction import DocExtraction
+from .download import Download
+from .indico_wrapper import IndicoWrapper
+from .reviewer import Reviewer
+from .workflow import Workflow
+
+__all__ = (
+    "Datasets",
+    "DocExtraction",
+    "Download",
+    "IndicoWrapper",
+    "Reviewer",
+    "Workflow",
+)

indico_toolkit/indico_wrapper/dataset.py

@@ -0,0 +1,114 @@
+from typing import List
+
+from indico import IndicoClient
+from indico.queries import (
+    AddDataToWorkflow,
+    AddFiles,
+    CreateDataset,
+    CreateEmptyDataset,
+    DeleteDataset,
+    GetDataset,
+)
+from indico.types import Dataset, OcrEngine, Workflow
+
+from .indico_wrapper import IndicoWrapper
+
+
+class Datasets(IndicoWrapper):
+    def __init__(self, client: IndicoClient):
+        self.client = client
+
+    def get_dataset(self, dataset_id: int):
+        return self.client.call(GetDataset(dataset_id))
+
+    def add_files_to_dataset(self, dataset_id: int, filepaths: List[str]) -> Dataset:
+        """
+        Upload documents to an existing dataset and wait for them to OCR
+        """
+        dataset = self.client.call(
+            AddFiles(
+                dataset_id=dataset_id, files=filepaths, autoprocess=True, wait=True
+            )
+        )
+        return dataset
+
+    def add_new_files_to_task(self, workflow_id: id, wait: bool = True) -> Workflow:
+        """
+        Add newly uploaded documents to an existing teach task given the task's
+        associated workflow ID
+
+        Args:
+            workflow_id (id): workflow ID associated with teach task
+            wait (bool, optional): wait for data to be added. Defaults to True.
+        """
+        workflow = self.client.call(AddDataToWorkflow(workflow_id, wait))
+        if wait:
+            print(f"Data added to all teach tasks associated with {workflow.id}")
+        return workflow
+
+    def create_empty_dataset(
+        self,
+        dataset_name: str,
+        dataset_type: str = "DOCUMENT",
+        ocr_engine: OcrEngine = OcrEngine.READAPI,
+    ) -> Dataset:
+        """
+        Create an empty dataset
+        Args:
+            name (str): Name of the dataset
+            dataset_type (str, optional): TEXT, IMAGE, or DOCUMENT.
+                Defaults to "DOCUMENT".
+        """
+        return self.client.call(
+            CreateEmptyDataset(dataset_name, dataset_type, ocr_engine)
+        )
+
+    def create_dataset(
+        self,
+        filepaths: List[str],
+        dataset_name: str,
+        ocr_engine: OcrEngine = OcrEngine.READAPI,
+    ) -> Dataset:
+        dataset = self.client.call(
+            CreateDataset(
+                name=dataset_name,
+                files=filepaths,
+                ocr_engine=ocr_engine,
+            )
+        )
+        self.dataset_id = dataset.id
+        return dataset
+
+    def delete_dataset(self, dataset_id: int) -> bool:
+        """
+        Returns True if operation is succesful
+        """
+        return self.client.call(DeleteDataset(id=dataset_id))
+
+    def get_dataset_metadata(self, dataset_id: int) -> List[dict]:
+        """
+        Get list of dataset files with information like file name, status, and number of
+        pages
+        """
+        query = """
+            query GetDataset($id: Int) {
+                dataset(id: $id) {
+                    id
+                    name
+                    files {
+                        id
+                        name
+                        numPages
+                        status
+                    }
+                }
+            }
+        """
+        dataset = self.graphQL_request(
+            graphql_query=query, variables={"id": dataset_id}
+        )
+        return dataset["dataset"]["files"]
+
+    def get_col_name_by_id(self, dataset_id: int, col_id: int) -> str:
+        dataset = self.get_dataset(dataset_id)
+        return next(c.name for c in dataset.datacolumns if c.id == col_id)

indico_toolkit/indico_wrapper/doc_extraction.py

@@ -0,0 +1,79 @@
+from typing import List, Union
+
+from indico import IndicoClient
+from indico.queries import DocumentExtraction, Job
+
+from ..ocr import CustomOcr, OnDoc, StandardOcr
+from .indico_wrapper import IndicoWrapper
+
+
+class DocExtraction(IndicoWrapper):
+    """
+    Class to support DocumentExtraction-related API calls
+    """
+
+    def __init__(
+        self,
+        client: IndicoClient,
+        preset_config: str = "standard",
+        custom_config: dict = None,
+    ):
+        """
+        Args:
+            preset_config (str): Options are simple, legacy, detailed, ondocument, and
+                standard.
+        """
+        self._preset_config = preset_config
+        self.client = client
+        self.json_config = {"preset_config": preset_config}
+        if custom_config:
+            self.json_config = custom_config
+
+    def run_ocr(
+        self, filepaths: List[str], text_setting: str = None
+    ) -> List[Union[StandardOcr, OnDoc, CustomOcr, str]]:
+        """
+        Args:
+            filepaths (List[str]): List of paths to local documents you would like to
+                submit for extraction
+            text_setting (str): Options are full_text and page_texts.
+
+        Returns:
+            extracted_data (List[Union[StandardOcr, OnDoc, CustomOcr, str]]): data from
+                DocumentExtraction converted to OCR objects or string text
+        """
+        jobs = self._submit_to_ocr(filepaths)
+        extracted_data = []
+        for ind, job in enumerate(jobs):
+            status = self.get_job_status(job.id, True)
+            if status.status == "SUCCESS":
+                result = self.get_storage_object(status.result)
+                if text_setting == "full_text":
+                    extracted_data.append(self._convert_ocr_objects(result).full_text)
+                elif text_setting == "page_texts":
+                    extracted_data.append(self._convert_ocr_objects(result).page_texts)
+                else:
+                    extracted_data.append(self._convert_ocr_objects(result))
+            else:
+                raise RuntimeError(
+                    f"{filepaths[ind]} {status.status}: {status.result}."
+                )
+        return extracted_data
+
+    def _submit_to_ocr(self, filepaths: List[str]) -> List[Job]:
+        return self.client.call(
+            DocumentExtraction(files=filepaths, json_config=self.json_config)
+        )
+
+    def _convert_ocr_objects(
+        self, extracted_data: Union[List[dict], dict]
+    ) -> Union[StandardOcr, OnDoc, CustomOcr]:
+        if self.json_config == {"preset_config": "ondocument"}:
+            return OnDoc(extracted_data)
+        elif (
+            self.json_config == {"preset_config": "standard"}
+            or self.json_config is None
+        ):
+            return StandardOcr(extracted_data)
+        else:
+            return CustomOcr(extracted_data)