"""Annotation publication, catalog coordination, and immutable set access."""
from __future__ import annotations
import bisect
import hashlib
import re
from collections.abc import Iterator, Mapping, Sequence
from types import MappingProxyType
from typing import Any, Protocol, overload
from signal_dataset._internal.json import (
frozen_mapping,
strict_json_loads,
)
from signal_dataset.annotation.codec import SafeTensorsAnnotationCodec
from signal_dataset.annotation.model import (
SET_NAME as _SET_NAME,
)
from signal_dataset.annotation.model import (
AnnotationRoot,
AnnotationShard,
)
from signal_dataset.config import StorageOptions
from signal_dataset.errors import CorruptDatasetError
from signal_dataset.record.model import (
AnnotationRecord,
RecordMetadata,
)
from signal_dataset.storage.contracts import ObjectStore, ShardStore
from signal_dataset.storage.paths import (
join_uri,
safe_relative_path,
)
from signal_dataset.storage.reference import ObjectReference
CATALOG = "annotations/catalog.json"
class DatasetContext(Protocol):
@property
def root(self) -> str: ...
@property
def snapshot_id(self) -> str: ...
@property
def record_metadata(self) -> Sequence[RecordMetadata]: ...
@property
def object_store(self) -> ObjectStore: ...
@property
def shard_store(self) -> ShardStore: ...
@property
def storage_options(self) -> StorageOptions: ...
def __len__(self) -> int: ...
def _parse(payload: bytes, options: StorageOptions, subject: str) -> dict[str, Any]:
if len(payload) > options.control.max_document_bytes:
raise CorruptDatasetError(f"{subject} exceeds max_document_bytes")
try:
value = strict_json_loads(
payload,
max_depth=options.control.max_json_depth,
max_nodes=options.control.max_json_nodes,
)
except (UnicodeDecodeError, ValueError) as exc:
raise CorruptDatasetError(f"invalid {subject}: {exc}") from exc
if not isinstance(value, dict):
raise CorruptDatasetError(f"{subject} must be an object")
return value
def _optional_generation(value: Any, subject: str) -> int | None:
if value is None:
return None
if isinstance(value, bool) or not isinstance(value, int) or value < 0:
raise CorruptDatasetError(f"{subject} generation must be a nonnegative integer")
return value
def _catalog_pointer(value: Mapping[str, Any]) -> tuple[str, int | None]:
try:
if set(value) != {"revision"} or not isinstance(value["revision"], dict):
raise TypeError("expected one revision object")
revision = value["revision"]
if set(revision) - {"path", "generation"} or "path" not in revision:
raise TypeError("unexpected revision reference fields")
path = revision["path"]
if not isinstance(path, str) or not re.fullmatch(
r"annotations/catalogs/[0-9a-f]{64}\.json", path
):
raise TypeError("invalid revision path")
return path, _optional_generation(revision.get("generation"), "revision")
except (KeyError, TypeError, ValueError) as exc:
if isinstance(exc, CorruptDatasetError):
raise
raise CorruptDatasetError(f"invalid annotation catalog pointer: {exc}") from exc
def _set_reference(name: str, value: Any) -> tuple[str, int | None, int]:
try:
if _SET_NAME.fullmatch(name) is None or not isinstance(value, dict):
raise TypeError("invalid set name or reference")
if set(value) - {"path", "generation", "stored_bytes"} or not {
"path",
"stored_bytes",
} <= set(value):
raise TypeError("unexpected set reference fields")
path = value["path"]
if not isinstance(path, str):
raise TypeError("invalid annotation root path")
path = safe_relative_path(path)
parts = path.split("/")
if (
len(parts) != 5
or parts[:3] != ["annotations", "sets", name]
or _SET_NAME.fullmatch(parts[3]) is None
or parts[4] != "root.json"
):
raise TypeError("invalid annotation root path")
size = value["stored_bytes"]
if isinstance(size, bool) or not isinstance(size, int) or size < 0:
raise TypeError("stored_bytes must be a nonnegative integer")
return path, _optional_generation(value.get("generation"), "annotation root"), size
except (TypeError, ValueError) as exc:
if isinstance(exc, CorruptDatasetError):
raise
raise CorruptDatasetError(f"invalid annotation set reference: {exc}") from exc
def _catalog_revision(value: Mapping[str, Any]) -> dict[str, Any]:
if set(value) != {"sets"} or not isinstance(value["sets"], dict):
raise CorruptDatasetError("annotation catalog revision must contain one sets object")
for name, reference in value["sets"].items():
_set_reference(name, reference)
return dict(value)
[docs]
class AnnotationSet(Sequence[AnnotationRecord]):
def __init__(
self,
root: str,
specification: AnnotationRoot,
objects: ObjectStore,
shards: ShardStore,
options: StorageOptions,
source_metadata: Sequence[RecordMetadata],
) -> None:
self.root = root
self.name = specification.name
self.publication_id = specification.publication_id
self.source_snapshot_id = specification.source_snapshot_id
self.metadata = frozen_mapping(specification.metadata)
self._specification = specification
self._objects = objects
self._shards = shards
self._options = options
self._codec = SafeTensorsAnnotationCodec(options=options)
self._starts = tuple(item.first_ordinal for item in specification.shards)
self._source_metadata = source_metadata
def __len__(self) -> int:
return self._specification.record_count
@overload
def __getitem__(self, index: int) -> AnnotationRecord: ...
@overload
def __getitem__(self, index: slice) -> list[AnnotationRecord]: ...
def __getitem__(self, index: int | slice) -> AnnotationRecord | list[AnnotationRecord]:
if isinstance(index, slice):
return [self[item] for item in range(*index.indices(len(self)))]
normalized = index + len(self) if index < 0 else index
if normalized < 0 or normalized >= len(self):
raise IndexError("annotation index out of range")
position = bisect.bisect_right(self._starts, normalized) - 1
shard = self._specification.shards[position]
uri = join_uri(self.root, shard.object.path)
info = self._objects.info(uri, generation=shard.object.generation)
if info.size != shard.object.stored_bytes:
raise CorruptDatasetError("annotation shard size does not match root")
payload = self._shards.read(
uri,
normalized - shard.first_ordinal,
generation=shard.object.generation,
options=self._options.reader_options,
file_reader_buffer_size=self._options.file_reader_buffer_size,
)
record = self._codec.decode(payload)
expected = self._source_metadata[normalized]
if record.source_index != normalized or record.source_record_id != expected.id:
raise CorruptDatasetError("annotation row identity does not match source ordinal")
return record
[docs]
class AnnotationCatalog(Mapping[str, AnnotationSet]):
def __init__(self, values: Mapping[str, AnnotationSet], revision: str | None) -> None:
self._values = MappingProxyType(dict(values))
self.revision = revision
def __getitem__(self, name: str) -> AnnotationSet:
return self._values[name]
def __iter__(self) -> Iterator[str]:
return iter(self._values)
def __len__(self) -> int:
return len(self._values)
def open_catalog(dataset: DatasetContext) -> AnnotationCatalog:
objects = dataset.object_store
try:
observed = objects.read_version(join_uri(dataset.root, CATALOG))
except FileNotFoundError:
return AnnotationCatalog({}, None)
options = dataset.storage_options
try:
pointer = _parse(observed.data, options, "annotation catalog pointer")
revision_path, revision_generation = _catalog_pointer(pointer)
except (KeyError, TypeError, AttributeError, ValueError) as exc:
if isinstance(exc, CorruptDatasetError):
raise
raise CorruptDatasetError(f"invalid annotation catalog pointer: {exc}") from exc
revision_payload = objects.read(
join_uri(dataset.root, revision_path),
generation=revision_generation,
)
revision_id = revision_path.rsplit("/", 1)[-1].removesuffix(".json")
if hashlib.sha256(revision_payload).hexdigest() != revision_id:
raise CorruptDatasetError("annotation catalog revision hash does not match payload")
document = _catalog_revision(_parse(revision_payload, options, "annotation catalog revision"))
values: dict[str, AnnotationSet] = {}
try:
set_references = document["sets"]
if not isinstance(set_references, dict):
raise CorruptDatasetError("annotation catalog sets must be an object")
except (KeyError, TypeError, AttributeError, ValueError) as exc:
if isinstance(exc, CorruptDatasetError):
raise
raise CorruptDatasetError(f"invalid annotation catalog revision: {exc}") from exc
for name, reference in set_references.items():
try:
path, root_generation, stored_bytes = _set_reference(name, reference)
root_uri = join_uri(dataset.root, path)
info = objects.info(root_uri, generation=root_generation)
if info.size != stored_bytes:
raise CorruptDatasetError("annotation root size does not match catalog")
root_document = _parse(
objects.read(root_uri, generation=root_generation),
options,
"annotation root",
)
except (KeyError, TypeError, AttributeError, ValueError) as exc:
if isinstance(exc, CorruptDatasetError):
raise
raise CorruptDatasetError(f"invalid annotation set reference: {exc}") from exc
try:
allowed_root = {
"name",
"publication_id",
"source_snapshot_id",
"record_count",
"metadata",
"shards",
}
if set(root_document) != allowed_root:
raise CorruptDatasetError("annotation root has unexpected fields")
shard_values = root_document["shards"]
if not isinstance(shard_values, list):
raise CorruptDatasetError("annotation root shards must be an array")
specification = AnnotationRoot(
root_document["name"],
root_document["publication_id"],
root_document["source_snapshot_id"],
root_document["record_count"],
tuple(
AnnotationShard(
item["first_ordinal"],
item["record_count"],
ObjectReference.from_dict(item["object"]),
)
for item in shard_values
),
root_document.get("metadata", {}),
)
except (KeyError, TypeError, AttributeError, ValueError) as exc:
if isinstance(exc, CorruptDatasetError):
raise
raise CorruptDatasetError(f"invalid annotation root: {exc}") from exc
expected = 0
for shard in specification.shards:
if shard.first_ordinal != expected:
raise CorruptDatasetError("annotation shard ranges are not contiguous")
expected += shard.record_count
if expected != specification.record_count:
raise CorruptDatasetError("annotation shard ranges do not match record count")
if specification.record_count != len(dataset):
raise CorruptDatasetError("annotation set is not dense for the source dataset")
if specification.name != name:
raise CorruptDatasetError("annotation catalog name does not match set root")
if specification.source_snapshot_id != dataset.snapshot_id:
raise CorruptDatasetError("annotation set targets another dataset snapshot")
values[name] = AnnotationSet(
dataset.root,
specification,
objects,
dataset.shard_store,
dataset.storage_options,
dataset.record_metadata,
)
return AnnotationCatalog(values, revision_path)