"""Annotation wire and publication models."""
from __future__ import annotations
import re
from collections.abc import Mapping
from dataclasses import dataclass, field
from typing import Any
from signal_dataset._internal.json import Json, frozen_mapping
from signal_dataset.errors import CorruptDatasetError
from signal_dataset.record.model import AnnotationRecord, AnnotationStatus
from signal_dataset.storage.reference import ObjectReference
SET_NAME = re.compile(r"[A-Za-z0-9][A-Za-z0-9._-]{0,127}")
@dataclass(frozen=True, slots=True)
class AnnotationShard:
first_ordinal: int
record_count: int
object: ObjectReference
def __post_init__(self) -> None:
if (
isinstance(self.first_ordinal, bool)
or not isinstance(self.first_ordinal, int)
or self.first_ordinal < 0
):
raise ValueError("annotation shard first_ordinal must be nonnegative")
if (
isinstance(self.record_count, bool)
or not isinstance(self.record_count, int)
or self.record_count < 1
):
raise ValueError("annotation shard record_count must be positive")
[docs]
@dataclass(frozen=True, slots=True)
class PublishedAnnotationShard:
name: str
publication_id: str
source_snapshot_id: str
logical_shard_id: str
first_ordinal: int
record_ids: tuple[str, ...]
uri: str
stored_bytes: int
generation: int | None
statuses: tuple[str, ...]
terminal_statuses: tuple[str, ...]
def published_shard_to_dict(value: PublishedAnnotationShard) -> dict[str, Any]:
return {
"name": value.name,
"publication_id": value.publication_id,
"source_snapshot_id": value.source_snapshot_id,
"logical_shard_id": value.logical_shard_id,
"first_ordinal": value.first_ordinal,
"record_ids": list(value.record_ids),
"uri": value.uri,
"stored_bytes": value.stored_bytes,
"generation": value.generation,
"statuses": list(value.statuses),
"terminal_statuses": list(value.terminal_statuses),
}
def published_shard_from_dict(value: Mapping[str, Any]) -> PublishedAnnotationShard:
try:
required = {
"name",
"publication_id",
"source_snapshot_id",
"logical_shard_id",
"first_ordinal",
"record_ids",
"uri",
"stored_bytes",
"generation",
"statuses",
"terminal_statuses",
}
if set(value) != required:
raise TypeError("unexpected pointer fields")
strings = tuple(
value[name]
for name in ("name", "publication_id", "source_snapshot_id", "logical_shard_id", "uri")
)
if any(not isinstance(item, str) or not item for item in strings):
raise TypeError("pointer string is invalid")
first, size, generation = value["first_ordinal"], value["stored_bytes"], value["generation"]
record_ids, statuses = value["record_ids"], value["statuses"]
terminal = value["terminal_statuses"]
if (
isinstance(first, bool)
or not isinstance(first, int)
or first < 0
or isinstance(size, bool)
or not isinstance(size, int)
or size < 0
or (
generation is not None
and (isinstance(generation, bool) or not isinstance(generation, int))
)
or not isinstance(record_ids, list)
or not record_ids
or any(not isinstance(item, str) or not item for item in record_ids)
or not isinstance(statuses, list)
or len(statuses) != len(record_ids)
or any(item not in {"success", "skipped", "failed"} for item in statuses)
or not isinstance(terminal, list)
or len(terminal) != len(record_ids)
or any(not isinstance(item, str) or not item for item in terminal)
):
raise TypeError("pointer scalar is invalid")
return PublishedAnnotationShard(
strings[0],
strings[1],
strings[2],
strings[3],
first,
tuple(record_ids),
strings[4],
size,
generation,
tuple(statuses),
tuple(terminal),
)
except (KeyError, TypeError, ValueError) as exc:
raise CorruptDatasetError(f"invalid annotation shard pointer: {exc}") from exc
def validated_published_shard(value: PublishedAnnotationShard) -> PublishedAnnotationShard:
return published_shard_from_dict(published_shard_to_dict(value))
@dataclass(frozen=True, slots=True)
class AnnotationRoot:
name: str
publication_id: str
source_snapshot_id: str
record_count: int
shards: tuple[AnnotationShard, ...]
metadata: Mapping[str, Json] = field(default_factory=dict)
def __post_init__(self) -> None:
if SET_NAME.fullmatch(self.name) is None or SET_NAME.fullmatch(self.publication_id) is None:
raise ValueError("annotation root identifiers must be portable")
if not self.source_snapshot_id:
raise ValueError("annotation source_snapshot_id must be non-empty")
if (
isinstance(self.record_count, bool)
or not isinstance(self.record_count, int)
or self.record_count < 0
):
raise ValueError("annotation record_count must be nonnegative")
object.__setattr__(self, "shards", tuple(self.shards))
object.__setattr__(self, "metadata", frozen_mapping(self.metadata))
__all__ = [
"AnnotationRecord",
"AnnotationRoot",
"AnnotationShard",
"AnnotationStatus",
"PublishedAnnotationShard",
]