"""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 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) 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))