unstructured-io--unstructured/unstructured/utils.py

from __future__ import annotations

import asyncio
import contextlib
import importlib
import inspect
import json
import os
import platform
import subprocess
import tempfile
import threading
from functools import lru_cache, wraps
from itertools import combinations
from typing import (
    TYPE_CHECKING,
    Any,
    Callable,
    Iterable,
    Iterator,
    List,
    Optional,
    Tuple,
    TypeVar,
    cast,
)

import requests
from typing_extensions import ParamSpec, TypeAlias

from unstructured.__version__ import __version__

if TYPE_CHECKING:
    from unstructured.documents.elements import Element, Text

# Box format: [x_bottom_left, y_bottom_left, x_top_right, y_top_right]
Box: TypeAlias = Tuple[float, float, float, float]
Point: TypeAlias = Tuple[float, float]
Points: TypeAlias = Tuple[Point, ...]

DATE_FORMATS = ("%Y-%m-%d", "%Y-%m-%dT%H:%M:%S", "%Y-%m-%d+%H:%M:%S", "%Y-%m-%dT%H:%M:%S%z")

_T = TypeVar("_T")
_P = ParamSpec("_P")


def get_call_args_applying_defaults(
    func: Callable[_P, List[Element]],
    *args: _P.args,
    **kwargs: _P.kwargs,
) -> dict[str, Any]:
    """Map both explicit and default arguments of decorated func call by param name."""
    sig = inspect.signature(func)
    call_args: dict[str, Any] = dict(**dict(zip(sig.parameters, args)), **kwargs)
    for arg in sig.parameters.values():
        if arg.name not in call_args and arg.default is not arg.empty:
            call_args[arg.name] = arg.default
    return call_args


def is_temp_file_path(file_path: str) -> bool:
    """True when file_path is in the Python-defined tempdir.

    The Python-defined temp directory is platform dependent (macOS != Linux != Windows)
    and can also be determined by an environment variable (TMPDIR, TEMP, or TMP).
    """
    return file_path.startswith(tempfile.gettempdir())


def save_as_jsonl(data: list[dict[str, Any]], filename: str) -> None:
    with open(filename, "w+") as output_file:
        output_file.writelines(json.dumps(datum) + "\n" for datum in data)


def read_from_jsonl(filename: str) -> list[dict[str, Any]]:
    with open(filename) as input_file:
        return [json.loads(line) for line in input_file]


def requires_dependencies(
    dependencies: str | list[str],
    extras: Optional[str] = None,
) -> Callable[[Callable[_P, _T]], Callable[_P, _T]]:
    if isinstance(dependencies, str):
        dependencies = [dependencies]

    def decorator(func: Callable[_P, _T]) -> Callable[_P, _T]:
        def run_check():
            missing_deps: List[str] = []
            for dep in dependencies:
                if not dependency_exists(dep):
                    missing_deps.append(dep)
            if len(missing_deps) > 0:
                raise ImportError(
                    f"Following dependencies are missing: {', '.join(missing_deps)}. "
                    + (
                        f"""Please install them using `pip install "unstructured[{extras}]"`."""
                        if extras
                        else f"Please install them using `pip install {' '.join(missing_deps)}`."
                    ),
                )

        @wraps(func)
        def wrapper(*args: _P.args, **kwargs: _P.kwargs):
            run_check()
            return func(*args, **kwargs)

        @wraps(func)
        async def wrapper_async(*args: _P.args, **kwargs: _P.kwargs):
            run_check()
            return await func(*args, **kwargs)

        if asyncio.iscoroutinefunction(func):
            return wrapper_async
        return wrapper

    return decorator


@lru_cache(maxsize=128)
def dependency_exists(dependency: str):
    try:
        importlib.import_module(dependency)
    except ImportError as e:
        # Check to make sure this isn't some unrelated import error.
        if dependency in repr(e):
            return False
    return True


def _first_and_remaining_iterator(it: Iterable[_T]) -> Tuple[_T, Iterator[_T]]:
    iterator = iter(it)
    try:
        out = next(iterator)
    except StopIteration:
        raise ValueError(
            "Expected at least 1 element in iterable from which to retrieve first, got empty "
            "iterable.",
        )
    return out, iterator


def first(it: Iterable[_T]) -> _T:
    """Returns the first item from an iterable. Raises an error if the iterable is empty."""
    out, _ = _first_and_remaining_iterator(it)
    return out


def only(it: Iterable[Any]) -> Any:
    """Returns the only element from a singleton iterable.

    Raises an error if the iterable is not a singleton.
    """
    out, iterator = _first_and_remaining_iterator(it)
    if any(True for _ in iterator):
        raise ValueError(
            "Expected only 1 element in passed argument, instead there are at least 2 elements.",
        )
    return out


def _telemetry_opt_out() -> bool:
    """True if telemetry should be disabled via env.

    DO_NOT_TRACK and SCARF_NO_ANALYTICS both follow the same rule: any non-empty
    value (after strip) opts out. See README/CHANGELOG for the public contract.
    """
    return bool((os.getenv("DO_NOT_TRACK") or "").strip()) or bool(
        (os.getenv("SCARF_NO_ANALYTICS") or "").strip()
    )


def _telemetry_opt_in() -> bool:
    """True if telemetry is explicitly enabled via env. Only 'true' and '1' opt in."""
    return (os.getenv("UNSTRUCTURED_TELEMETRY_ENABLED") or "").strip().lower() in (
        "true",
        "1",
    )


def scarf_analytics():
    """Send a lightweight analytics ping. Off by default.

    Set UNSTRUCTURED_TELEMETRY_ENABLED=true to opt in.
    Opt-out env vars (DO_NOT_TRACK, SCARF_NO_ANALYTICS): any non-empty value opts out.
    """
    if _telemetry_opt_out() or not _telemetry_opt_in():
        return

    try:
        subprocess.check_output(["nvidia-smi"], stderr=subprocess.DEVNULL)
        gpu_present = True
    except (OSError, subprocess.CalledProcessError):
        gpu_present = False

    python_version = ".".join(platform.python_version().split(".")[:2])

    with contextlib.suppress(Exception):
        requests.get(
            "https://packages.unstructured.io/python-telemetry",
            params={
                "version": __version__,
                "platform": platform.system(),
                "python": python_version,
                "arch": platform.machine(),
                "gpu": str(gpu_present),
                "dev": str("dev" in __version__).lower(),
            },
            timeout=10,
        )


def ngrams(s: list[str], n: int) -> list[tuple[str, ...]]:
    """Generate n-grams from a list of strings where `n` (int) is the size of each n-gram."""

    if n <= 0:
        raise ValueError(f"n must be positive, received n = {n}")
    return [tuple(s[i : i + n]) for i in range(len(s) - n + 1)]


def calculate_shared_ngram_percentage(
    first_string: str,
    second_string: str,
    n: int,
) -> tuple[float, set[tuple[str, ...]]]:
    """Calculate the percentage of common_ngrams between string A and B with reference to A"""
    if not n:
        return 0, set()
    first_string_ngrams = ngrams(first_string.split(), n)
    second_string_ngrams = ngrams(second_string.split(), n)

    if not first_string_ngrams:
        return 0, set()

    common_ngrams = set(first_string_ngrams) & set(second_string_ngrams)
    percentage = (len(common_ngrams) / len(first_string_ngrams)) * 100
    return percentage, common_ngrams


def calculate_largest_ngram_percentage(
    first_string: str, second_string: str
) -> tuple[float, set[tuple[str, ...]], str]:
    """From two strings, calculate the shared ngram percentage.

    Returns a tuple containing...
        - The largest n-gram percentage shared between the two strings.
        - A set containing the shared n-grams found during the calculation.
        - A string representation of the size of the largest shared n-grams found.
    """
    shared_ngrams: set[tuple[str, ...]] = set()
    if len(first_string.split()) < len(second_string.split()):
        n = len(first_string.split()) - 1
    else:
        n = len(second_string.split()) - 1
        first_string, second_string = second_string, first_string
    ngram_percentage = 0
    # Start from the biggest ngram possible (`n`) until the ngram_percentage is >0.0% or n == 0
    while not ngram_percentage:
        ngram_percentage, shared_ngrams = calculate_shared_ngram_percentage(
            first_string,
            second_string,
            n,
        )
        if n == 0:
            break
        else:
            n -= 1
    return round(ngram_percentage, 2), shared_ngrams, str(n + 1)


def is_parent_box(parent_target: Box, child_target: Box, add: float = 0.0) -> bool:
    """True if the child_target bounding box is nested in the parent_target.

    Box format: [x_bottom_left, y_bottom_left, x_top_right, y_top_right].
    The parameter 'add' is the pixel error tolerance for extra pixels outside the parent region
    """
    if len(parent_target) != 4:
        return False
    parent_targets = [0, 0, 0, 0]
    if add and len(parent_target) == 4:
        parent_targets = list(parent_target)
        parent_targets[0] -= add
        parent_targets[1] -= add
        parent_targets[2] += add
        parent_targets[3] += add

    if (
        len(child_target) == 4
        and (child_target[0] >= parent_targets[0] and child_target[1] >= parent_targets[1])
        and (child_target[2] <= parent_targets[2] and child_target[3] <= parent_targets[3])
    ):
        return True
    return len(child_target) == 2 and (
        parent_targets[0] <= child_target[0] <= parent_targets[2]
        and parent_targets[1] <= child_target[1] <= parent_targets[3]
    )


def calculate_overlap_percentage(
    box1: Points,
    box2: Points,
    intersection_ratio_method: str = "total",
) -> tuple[float, float, float, float]:
    """Calculate the percentage of overlapped region.

    Calculate the percentage with reference to
    the biggest element-region (intersection_ratio_method="parent"),
    the smallest element-region (intersection_ratio_method="partial"), or
    the disjunctive union region (intersection_ratio_method="total")
    """
    x1, y1 = box1[0]
    x2, y2 = box1[2]
    x3, y3 = box2[0]
    x4, y4 = box2[2]
    area_box1 = (x2 - x1) * (y2 - y1)
    area_box2 = (x4 - x3) * (y4 - y3)
    x_intersection1 = max(x1, x3)
    y_intersection1 = max(y1, y3)
    x_intersection2 = min(x2, x4)
    y_intersection2 = min(y2, y4)
    intersection_area = max(0, x_intersection2 - x_intersection1) * max(
        0,
        y_intersection2 - y_intersection1,
    )
    max_area = max(area_box1, area_box2)
    min_area = min(area_box1, area_box2)
    total_area = area_box1 + area_box2

    if intersection_ratio_method == "parent":
        if max_area == 0:
            return 0, 0, 0, 0
        overlap_percentage = (intersection_area / max_area) * 100

    elif intersection_ratio_method == "partial":
        if min_area == 0:
            return 0, 0, 0, 0
        overlap_percentage = (intersection_area / min_area) * 100

    else:
        if (area_box1 + area_box2) == 0:
            return 0, 0, 0, 0

        overlap_percentage = (intersection_area / (area_box1 + area_box2 - intersection_area)) * 100

    return round(overlap_percentage, 2), max_area, min_area, total_area


def identify_overlapping_case(
    box_pair: list[Points] | tuple[Points, Points],
    label_pair: list[str] | tuple[str, str],
    text_pair: list[str] | tuple[str, str],
    ix_pair: list[str] | tuple[str, str],
    sm_overlap_threshold: float = 10.0,
):
    """Classifies the overlapping case for an element_pair input.

    There are 5 cases of overlapping:
        'Small partial overlap'
        'Partial overlap with empty content'
        'Partial overlap with duplicate text (sharing 100% of the text)'
        'Partial overlap without sharing text'
        'Partial overlap sharing {calculate_largest_ngram_percentage(...)}% of the text'

    Returns:
    overlapping_elements: List[str] - List of element types with their `ix` value.
        Ex: ['Title(ix=0)']
    overlapping_case: str - See list of cases above
    overlap_percentage: float
    largest_ngram_percentage: float
    max_area: float
    min_area: float
    total_area: float
    """
    overlapping_elements, overlapping_case, overlap_percentage, largest_ngram_percentage = (
        None,
        None,
        None,
        None,
    )
    box1, box2 = box_pair
    type1, type2 = label_pair
    text1, text2 = text_pair
    ix_element1, ix_element2 = ix_pair
    overlap_percentage, max_area, min_area, total_area = calculate_overlap_percentage(
        box1,
        box2,
        intersection_ratio_method="partial",
    )
    if overlap_percentage < sm_overlap_threshold:
        overlapping_elements = [
            f"{type1}(ix={ix_element1})",
            f"{type2}(ix={ix_element2})",
        ]
        overlapping_case = "Small partial overlap"

    else:
        if not text1:
            overlapping_elements = [
                f"{type1}(ix={ix_element1})",
                f"{type2}(ix={ix_element2})",
            ]
            overlapping_case = f"partial overlap with empty content in {type1}"

        elif not text2:
            overlapping_elements = [
                f"{type2}(ix={ix_element2})",
                f"{type1}(ix={ix_element1})",
            ]
            overlapping_case = f"partial overlap with empty content in {type2}"

        elif text1 in text2 or text2 in text1:
            overlapping_elements = [
                f"{type1}(ix={ix_element1})",
                f"{type2}(ix={ix_element2})",
            ]
            overlapping_case = "partial overlap with duplicate text"

        else:
            largest_ngram_percentage, _, largest_n = calculate_largest_ngram_percentage(
                text1, text2
            )
            largest_ngram_percentage = round(largest_ngram_percentage, 2)
            if not largest_ngram_percentage:
                overlapping_elements = [
                    f"{type1}(ix={ix_element1})",
                    f"{type2}(ix={ix_element2})",
                ]
                overlapping_case = "partial overlap without sharing text"

            else:
                overlapping_elements = [
                    f"{type1}(ix={ix_element1})",
                    f"{type2}(ix={ix_element2})",
                ]
                ref_type = type1 if len(text1.split()) < len(text2.split()) else type2
                ref_type = "of the text from" + ref_type + f"({largest_n}-gram)"
                overlapping_case = f"partial overlap sharing {largest_ngram_percentage}% {ref_type}"
    return (
        overlapping_elements,
        overlapping_case,
        overlap_percentage,
        largest_ngram_percentage,
        max_area,
        min_area,
        total_area,
    )


def _convert_coordinates_to_box(coordinates: Points):
    """Accepts a set of Points and returns the lower-left and upper-right coordinates.

    Expects four coordinates representing the corners of a rectangle, listed in this order:
    bottom-left, top-left, top-right, bottom-right.
    """
    x_bottom_left_1, y_bottom_left_1 = coordinates[0]
    x_top_right_1, y_top_right_1 = coordinates[2]
    return x_bottom_left_1, y_bottom_left_1, x_top_right_1, y_top_right_1


# x1, y1 = box1[0]
def identify_overlapping_or_nesting_case(
    box_pair: list[Points] | tuple[Points, Points],
    label_pair: list[str] | tuple[str, str],
    text_pair: list[str] | tuple[str, str],
    nested_error_tolerance_px: int = 5,
    sm_overlap_threshold: float = 10.0,
):
    """Identify if overlapping or nesting elements exist and, if so, the type of overlapping case.

    Returns:
    overlapping_elements: List[str] - List of element types & their `ix` value. Ex: ['Title(ix=0)']
    overlapping_case: str - See list of cases above
    overlap_percentage: float
    overlap_percentage_total: float
    largest_ngram_percentage: float
    max_area: float
    min_area: float
    total_area: float
    """
    box1, box2 = box_pair
    type1, type2 = label_pair
    ix_element1 = "".join([ch for ch in type1 if ch.isnumeric()])
    ix_element2 = "".join([ch for ch in type2 if ch.isnumeric()])
    type1 = type1[3:].strip()
    type2 = type2[3:].strip()
    box1_corners = _convert_coordinates_to_box(box1)
    box2_corners = _convert_coordinates_to_box(box2)
    x_bottom_left_1, y_bottom_left_1, x_top_right_1, y_top_right_1 = box1_corners
    x_bottom_left_2, y_bottom_left_2, x_top_right_2, y_top_right_2 = box2_corners

    horizontal_overlap = x_bottom_left_1 < x_top_right_2 and x_top_right_1 > x_bottom_left_2
    vertical_overlap = y_bottom_left_1 < y_top_right_2 and y_top_right_1 > y_bottom_left_2
    (
        overlapping_elements,
        parent_element,
        overlapping_case,
        overlap_percentage,
        overlap_percentage_total,
        largest_ngram_percentage,
    ) = (
        None,
        None,
        None,
        None,
        None,
        None,
    )
    max_area, min_area, total_area = None, None, None

    if horizontal_overlap and vertical_overlap:
        overlap_percentage_total, _, _, _ = calculate_overlap_percentage(
            box1,
            box2,
            intersection_ratio_method="total",
        )
        overlap_percentage, max_area, min_area, total_area = calculate_overlap_percentage(
            box1,
            box2,
            intersection_ratio_method="parent",
        )

        if is_parent_box(box1_corners, box2_corners, add=nested_error_tolerance_px):
            overlapping_elements = [
                f"{type1}(ix={ix_element1})",
                f"{type2}(ix={ix_element2})",
            ]
            overlapping_case = f"nested {type2} in {type1}"
            overlap_percentage = 100
            parent_element = f"{type1}(ix={ix_element1})"

        elif is_parent_box(box2_corners, box1_corners, add=nested_error_tolerance_px):
            overlapping_elements = [
                f"{type2}(ix={ix_element2})",
                f"{type1}(ix={ix_element1})",
            ]
            overlapping_case = f"nested {type1} in {type2}"
            overlap_percentage = 100
            parent_element = f"{type2}(ix={ix_element2})"

        else:
            (
                overlapping_elements,
                overlapping_case,
                overlap_percentage,
                largest_ngram_percentage,
                max_area,
                min_area,
                total_area,
            ) = identify_overlapping_case(
                box_pair,
                label_pair,
                text_pair,
                (ix_element1, ix_element2),
                sm_overlap_threshold=sm_overlap_threshold,
            )
    return (
        overlapping_elements,
        parent_element,
        overlapping_case,
        overlap_percentage or 0,
        overlap_percentage_total or 0,
        largest_ngram_percentage or 0,
        max_area or 0,
        min_area or 0,
        total_area or 0,
    )


def catch_overlapping_and_nested_bboxes(
    elements: list["Text"],
    nested_error_tolerance_px: int = 5,
    sm_overlap_threshold: float = 10.0,
) -> tuple[bool, list[dict[str, Any]]]:
    """Catch overlapping and nested bounding boxes cases across a list of elements."""

    num_pages = elements[-1].metadata.page_number or 0
    pages_of_bboxes: list[list[Points]] = [[] for _ in range(num_pages)]

    text_labels: list[list[str]] = [[] for _ in range(num_pages)]
    text_content: list[list[str]] = [[] for _ in range(num_pages)]

    for ix, element in enumerate(elements):
        page_number = element.metadata.page_number or 1
        n_page_to_ix = page_number - 1
        if element.metadata.coordinates:
            box = cast(Points, element.metadata.coordinates.to_dict()["points"])
            pages_of_bboxes[n_page_to_ix].append(box)
        text_labels[n_page_to_ix].append(f"{ix}. {element.category}")
        text_content[n_page_to_ix].append(element.text)

    document_with_overlapping_flag = False
    overlapping_cases: list[dict[str, Any]] = []
    for page_number, (page_bboxes, page_labels, page_text) in enumerate(
        zip(pages_of_bboxes, text_labels, text_content),
        start=1,
    ):
        page_bboxes_combinations = list(combinations(page_bboxes, 2))
        page_labels_combinations = list(combinations(page_labels, 2))
        text_content_combinations = list(combinations(page_text, 2))

        for box_pair, label_pair, text_pair in zip(
            page_bboxes_combinations,
            page_labels_combinations,
            text_content_combinations,
        ):
            (
                overlapping_elements,
                parent_element,
                overlapping_case,
                overlap_percentage,
                overlap_percentage_total,
                largest_ngram_percentage,
                max_area,
                min_area,
                total_area,
            ) = identify_overlapping_or_nesting_case(
                box_pair,
                label_pair,
                text_pair,
                nested_error_tolerance_px,
                sm_overlap_threshold,
            )

            if overlapping_case:
                overlapping_cases.append(
                    {
                        "overlapping_elements": overlapping_elements,
                        "parent_element": parent_element,
                        "overlapping_case": overlapping_case,
                        "overlap_percentage": f"{overlap_percentage}%",
                        "metadata": {
                            "largest_ngram_percentage": largest_ngram_percentage,
                            "overlap_percentage_total": f"{overlap_percentage_total}%",
                            "max_area": f"{round(max_area, 2)}pxˆ2",
                            "min_area": f"{round(min_area, 2)}pxˆ2",
                            "total_area": f"{round(total_area, 2)}pxˆ2",
                        },
                    },
                )
                document_with_overlapping_flag = True

    return document_with_overlapping_flag, overlapping_cases


def group_elements_by_parent_id(
    elements: Iterable["Element"],
    assign_orphans: bool = False,
) -> dict[Optional[str], list["Element"]]:
    """Group elements by their parent_id metadata field.

    Elements with the same parent_id are grouped together.

    Args:
        elements: An iterable of Element objects to group.
        assign_orphans: If True, elements with no parent_id (None) will be assigned to
            the same group as the previous element. If False (default), elements with
            no parent are grouped under the None key.

    Returns:
        A dictionary mapping parent_id values to lists of elements sharing that parent_id.

    Example:
        >>> elements = partition("example.pdf")
        >>> grouped = group_elements_by_parent_id(elements)
        >>> for parent_id, children in grouped.items():
        ...     print(f"Parent {parent_id}: {len(children)} children")

        >>> # Assign orphan elements to previous element's group
        >>> grouped = group_elements_by_parent_id(elements, assign_orphans=True)
    """
    from collections import defaultdict

    groups: dict[Optional[str], list["Element"]] = defaultdict(list)
    last_parent_id: Optional[str] = None

    for element in elements:
        parent_id = getattr(element.metadata, "parent_id", None)

        if parent_id is None and assign_orphans:
            parent_id = last_parent_id
        elif parent_id is not None:
            last_parent_id = parent_id

        groups[parent_id].append(element)

    return dict(groups)


class FileHandler:
    def __init__(self, file_path: str):
        self.file_path = file_path
        self.lock = threading.Lock()

    def read_file(self):
        with self.lock:
            with open(self.file_path) as file:
                data = file.read()
            return data

    def write_file(self, data: str) -> None:
        with self.lock:
            with open(self.file_path, "w") as file:
                file.write(data)

    def cleanup_file(self):
        with self.lock:
            if os.path.exists(self.file_path):
                os.remove(self.file_path)