"""Helpers used across multiple partitioners to compute metadata.""" from __future__ import annotations import copy import datetime as dt import functools import itertools import os from typing import Any, Callable, Iterator, Sequence from typing_extensions import ParamSpec from unstructured.documents.elements import Element, ElementMetadata 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", ], } 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) elements = list( apply_lang_metadata( elements=elements, languages=languages, detect_language_per_element=detect_language_per_element, ) ) # == 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. 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. """ # -- generate sequence number for each element on a page -- page_numbers = [e.metadata.page_number for e in elements] page_seq_numbers = [ seq_on_page for _, group in itertools.groupby(page_numbers) for seq_on_page, _ in enumerate(group) ] for element, seq_on_page_counter in zip(elements, page_seq_numbers): element.id_to_hash(seq_on_page_counter) 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))