Files
wehub-resource-sync 76d991c447
Auto Update PR / update-prs (push) Has been cancelled
CI / format-check (push) Has been cancelled
CI / test (3.10) (push) Has been cancelled
CI / test (3.11) (push) Has been cancelled
CI / test (3.12) (push) Has been cancelled
CI / live-api-tests (push) Has been cancelled
CI / plugin-integration-test (push) Has been cancelled
CI / ollama-integration-test (push) Has been cancelled
CI / test-fork-pr (push) Has been cancelled
chore: import upstream snapshot with attribution
2026-07-13 12:37:14 +08:00

125 lines
3.7 KiB
Python

# Copyright 2025 Google LLC.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Library for data conversion between AnnotatedDocument and JSON."""
from __future__ import annotations
import dataclasses
import enum
import numbers
from typing import Any, Iterable, Mapping
from langextract.core import data
from langextract.core import tokenizer
def enum_asdict_factory(items: Iterable[tuple[str, Any]]) -> dict[str, Any]:
"""Custom dict_factory for dataclasses.asdict.
Recursively converts dataclass instances, converts enum values to their
underlying values, converts integral numeric types to int, and skips any
field whose name starts with an underscore.
Args:
items: An iterable of (key, value) pairs from fields of a dataclass.
Returns:
A mapping of field names to their values, with special handling for
dataclasses, enums, and numeric types.
"""
result: dict[str, Any] = {}
for key, value in items:
# Skip internal fields.
if key.startswith("_"):
continue
if dataclasses.is_dataclass(value):
result[key] = dataclasses.asdict(value, dict_factory=enum_asdict_factory)
elif isinstance(value, enum.Enum):
result[key] = value.value
elif isinstance(value, numbers.Integral) and not isinstance(value, bool):
result[key] = int(value)
else:
result[key] = value
return result
def annotated_document_to_dict(
adoc: data.AnnotatedDocument | None,
) -> dict[str, Any]:
"""Converts an AnnotatedDocument into a Python dict.
This function converts an AnnotatedDocument object into a Python dict, making
it easier to serialize or deserialize the document. Enum values and NumPy
integers are converted to their underlying values, while other data types are
left unchanged. Private fields with an underscore prefix are not included in
the output.
Args:
adoc: The AnnotatedDocument object to convert.
Returns:
A Python dict representing the AnnotatedDocument.
"""
if not adoc:
return {}
result = dataclasses.asdict(adoc, dict_factory=enum_asdict_factory)
result["document_id"] = adoc.document_id
return result
def dict_to_annotated_document(
adoc_dic: Mapping[str, Any],
) -> data.AnnotatedDocument:
"""Converts a Python dict back to an AnnotatedDocument.
Args:
adoc_dic: A Python dict representing an AnnotatedDocument.
Returns:
An AnnotatedDocument object.
"""
if not adoc_dic:
return data.AnnotatedDocument()
for extractions in adoc_dic.get("extractions", []):
token_int = extractions.get("token_interval")
if token_int:
extractions["token_interval"] = tokenizer.TokenInterval(**token_int)
else:
extractions["token_interval"] = None
char_int = extractions.get("char_interval")
if char_int:
extractions["char_interval"] = data.CharInterval(**char_int)
else:
extractions["char_interval"] = None
status_str = extractions.get("alignment_status")
if status_str:
extractions["alignment_status"] = data.AlignmentStatus(status_str)
else:
extractions["alignment_status"] = None
return data.AnnotatedDocument(
document_id=adoc_dic.get("document_id"),
text=adoc_dic.get("text"),
extractions=[
data.Extraction(**ent) for ent in adoc_dic.get("extractions", [])
],
)