unstructured-io--unstructured/unstructured/partition/common/metadata.py

"""Helpers used across multiple partitioners to compute metadata."""

from __future__ import annotations

import copy
import datetime as dt
import functools
import os
from typing import Any, Callable, Iterator, Sequence

from typing_extensions import ParamSpec

from unstructured.documents.elements import Element, ElementMetadata, ListItem, Title
from unstructured.file_utils.model import FileType
from unstructured.partition.common.lang import apply_lang_metadata
from unstructured.utils import get_call_args_applying_defaults

_P = ParamSpec("_P")


def get_last_modified_date(filename: str) -> str | None:
    """Modification time of file at path `filename`, if it exists.

    Returns `None` when `filename` is not a path to a file on the local filesystem.

    Otherwise returns date and time in ISO 8601 string format (YYYY-MM-DDTHH:MM:SS) like
    "2024-03-05T17:02:53".
    """
    if not os.path.isfile(filename):
        return None

    modify_date = dt.datetime.fromtimestamp(os.path.getmtime(filename))
    return modify_date.strftime("%Y-%m-%dT%H:%M:%S%z")


HIERARCHY_RULE_SET = {
    "Title": [
        "Text",
        "UncategorizedText",
        "NarrativeText",
        "ListItem",
        "BulletedText",
        "Table",
        "FigureCaption",
        "CheckBox",
        "Table",
    ],
    "Header": [
        "Title",
        "Text",
        "UncategorizedText",
        "NarrativeText",
        "ListItem",
        "BulletedText",
        "Table",
        "FigureCaption",
        "CheckBox",
        "Table",
    ],
}


# Canonical HTML heading levels -> zero-indexed category_depth. The HTML spec
# defines exactly six heading levels (there is no h7), so this closed mapping is
# the single source of truth for both the depth value and the heading-tag set
# (HEADING_TAGS is derived from it, not a second copy).
_HEADING_DEPTH = {"h1": 0, "h2": 1, "h3": 2, "h4": 3, "h5": 4, "h6": 5}
HEADING_TAGS = tuple(_HEADING_DEPTH)


def category_depth_from_html_tag(
    ElementCls: type[Element], tag: str | None, list_ancestor_count: int = 0
) -> int | None:
    """Compute `category_depth` from an element's HTML heading level (not DOM-nesting depth).

    This is the canonical mapping used by both the v1 HTML parser and the v2 (ontology) HTML
    converter so the two paths agree on what `category_depth` means:

    - `Title` (which includes ontology Title/Subtitle/Heading, i.e. ``<h1>``-``<h6>``): the heading
      level, zero-indexed -- ``h1`` -> 0, ``h2`` -> 1, ... ``h6`` -> 5. A `Title` whose tag is not a
      heading (e.g. a styled paragraph) is treated as a top-level heading (0).
    - `ListItem`: the number of enclosing list containers (``ol``/``ul``/``dl``), passed in by the
      caller (the v1 HTML parser, which computes it from list nesting). The v2 converter serializes
      a whole ``ol``/``ul``/``dl`` as one element and never emits a standalone ``ListItem``, so it
      does not use this.
    - Everything else: ``None`` (no meaningful depth).

    `tag` is the element's HTML tag name (e.g. ``"h2"``); it may be ``None`` for derived elements.
    """
    if ElementCls is ListItem:
        return list_ancestor_count

    if ElementCls is Title:
        return _HEADING_DEPTH.get(tag, 0)

    return None


def set_element_hierarchy(
    elements: Sequence[Element], ruleset: dict[str, list[str]] = HIERARCHY_RULE_SET
) -> list[Element]:
    """Sets `.metadata.parent_id` for each element it applies to.

    `parent_id` assignment is based on the element's category and depth. The importance of an
    element's category is determined by a rule set. The rule set trumps category_depth. That is,
    category_depth is only relevant when elements are of the same category.
    """
    stack: list[Element] = []
    for element in elements:
        if element.metadata.parent_id is not None:
            continue
        parent_id = None
        element_category = getattr(element, "category", None)
        element_category_depth = getattr(element.metadata, "category_depth", 0) or 0

        # -- skip elements without a category --
        if not element_category:
            continue

        while stack:
            top_element: Element = stack[-1]
            top_element_category = getattr(top_element, "category")
            top_element_category_depth = (
                getattr(
                    top_element.metadata,
                    "category_depth",
                    0,
                )
                or 0
            )

            if (
                top_element_category == element_category
                and top_element_category_depth < element_category_depth
            ) or (
                top_element_category != element_category
                and element_category in ruleset.get(top_element_category, [])
            ):
                parent_id = top_element.id
                break

            stack.pop()

        element.metadata.parent_id = parent_id
        stack.append(element)

    return list(elements)


# ================================================================================================
# METADATA POST-PARTITIONING PROCESSING DECORATOR
# ================================================================================================


def apply_metadata(
    file_type: FileType | None = None,
) -> Callable[[Callable[_P, list[Element]]], Callable[_P, list[Element]]]:
    """Post-process element-metadata for this document.

    This decorator adds a post-processing step to a partitioner, primarily to apply metadata that
    is common to all partitioners. It assumes the following responsibilities:

      - Hash element-ids. Computes and applies SHA1 hash element.id when `unique_element_ids`
        argument is False.

      - Element Hierarchy. Computes and applies `parent_id` metadata based on `category_depth`
        etc. added by partitioner.

      - Language metadata. Computes and applies `language` metadata based on a language detection
        model.

      - Apply `filetype` (MIME-type) metadata. There are three cases; first one in this order that
        applies is used:

          - `metadata_file_type` argument is present in call, use that.
          - `file_type` decorator argument is populated, use that.
          - `file_type` decorator argument is omitted or None, don't apply `.metadata.filetype`
            (assume the partitioner will do that for itself, like `partition_image()`.

      - Replace `filename` with `metadata_filename` when present.

      - Replace `last_modified` with `metadata_last_modified` when present.

      - Apply `url` metadata when present.
    """

    def decorator(func: Callable[_P, list[Element]]) -> Callable[_P, list[Element]]:
        """The decorator function itself.

        This function is returned by the `apply_metadata()` function and is the actual decorator.
        Think of `apply_metadata()` as a factory function that configures this decorator, in
        particular by setting its `file_type` value.
        """

        @functools.wraps(func)
        def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> list[Element]:
            elements = func(*args, **kwargs)
            call_args = get_call_args_applying_defaults(func, *args, **kwargs)

            # ------------------------------------------------------------------------------------
            # unique-ify elements
            # ------------------------------------------------------------------------------------
            # Do this first to ensure all following operations behave as expected. It's easy for a
            # partitioner to re-use an element or metadata instance when its values are common to
            # multiple elements. This can lead to very hard-to diagnose bugs downstream when
            # mutating one element unexpectedly also mutates others (because they are the same
            # instance).
            # ------------------------------------------------------------------------------------

            elements = _uniqueify_elements_and_metadata(elements)

            # ------------------------------------------------------------------------------------
            # apply metadata - do this first because it affects the hash computation.
            # ------------------------------------------------------------------------------------

            # -- `language` - auto-detect language (e.g. eng, spa) --
            languages = call_args.get("languages")
            detect_language_per_element = call_args.get("detect_language_per_element", False)
            language_fallback = call_args.get("language_fallback")
            elements = list(
                apply_lang_metadata(
                    elements=elements,
                    languages=languages,
                    detect_language_per_element=detect_language_per_element,
                    language_fallback=language_fallback,
                )
            )

            # == apply filetype, filename, last_modified, and url metadata ===================
            metadata_kwargs: dict[str, Any] = {}

            # -- `filetype` (MIME-type) metadata --
            metadata_file_type = call_args.get("metadata_file_type") or file_type
            if metadata_file_type is not None:
                metadata_kwargs["filetype"] = metadata_file_type.mime_type

            # -- `filename` metadata - override with metadata_filename when it's present --
            filename = call_args.get("metadata_filename") or call_args.get("filename")
            if filename:
                metadata_kwargs["filename"] = filename

            # -- `last_modified` metadata - override with metadata_last_modified when present --
            metadata_last_modified = call_args.get("metadata_last_modified")
            if metadata_last_modified:
                metadata_kwargs["last_modified"] = metadata_last_modified

            # -- `url` metadata - record url when present --
            url = call_args.get("url")
            if url:
                metadata_kwargs["url"] = url

            # -- update element.metadata in single pass --
            for element in elements:
                # NOTE(robinson) - Attached files have already run through this logic in their own
                # partitioning function
                if element.metadata.attached_to_filename:
                    continue
                element.metadata.update(ElementMetadata(**metadata_kwargs))

            # ------------------------------------------------------------------------------------
            # compute hash ids (when so requestsd)
            # ------------------------------------------------------------------------------------

            # -- Compute and apply hash-ids if the user does not want UUIDs. Note this mutates the
            # -- elements themselves, not their metadata.
            unique_element_ids: bool = call_args.get("unique_element_ids", False)
            if unique_element_ids is False:
                elements = _assign_hash_ids(elements)

            # ------------------------------------------------------------------------------------
            # assign parent-id - do this after hash computation so parent-id is stable.
            # ------------------------------------------------------------------------------------

            # -- `parent_id` - process category-level etc. to assign parent-id --
            elements = set_element_hierarchy(elements)

            return elements

        return wrapper

    return decorator


def _assign_hash_ids(elements: list[Element]) -> list[Element]:
    """Converts `.id` of each element from UUID to hash and remaps `parent_id` accordingly.

    The hash is based on the `.text` of the element, but also on its page-number and sequence number
    on that page. This provides for deterministic results even when the document is split into one
    or more fragments for parallel processing.

    After hashing, any `element.metadata.parent_id` that references a known original UUID is
    updated to the corresponding new hash ID. Parent IDs that do not appear in the mapping (e.g.
    because the parent element was filtered out before hashing, or the ID was set manually to an
    external value) are left unchanged.
    """
    # -- generate sequence number for each element on a page --
    page_seq_counts = {}
    id_mapping = {}
    for element in elements:
        page_number = element.metadata.page_number
        seq_on_page_counter = page_seq_counts.get(page_number, 0)
        original_id = element.id
        element.id_to_hash(seq_on_page_counter)
        id_mapping[original_id] = element.id
        page_seq_counts[page_number] = seq_on_page_counter + 1

    for element in elements:
        if element.metadata.parent_id is not None and element.metadata.parent_id in id_mapping:
            element.metadata.parent_id = id_mapping[element.metadata.parent_id]

    return elements


def _uniqueify_elements_and_metadata(elements: list[Element]) -> list[Element]:
    """Ensure each of `elements` and their metadata are unique instances.

    This prevents hard-to-diagnose bugs downstream when mutating one element unexpectedly also
    mutates others because they are the same instance.
    """

    def iter_unique_elements(elements: list[Element]) -> Iterator[Element]:
        """Substitute deep-copies of any non-unique elements or metadata in `elements`."""
        seen_elements: set[int] = set()
        seen_metadata: set[int] = set()

        for element in elements:
            if id(element) in seen_elements:
                element = copy.deepcopy(element)
            if id(element.metadata) in seen_metadata:
                element.metadata = copy.deepcopy(element.metadata)
            seen_elements.add(id(element))
            seen_metadata.add(id(element.metadata))
            yield element

    return list(iter_unique_elements(elements))