108 lines
3.7 KiB
Python
108 lines
3.7 KiB
Python
"""Definition of AnnotationRepository."""
|
|
|
|
from __future__ import annotations
|
|
from uuid import UUID
|
|
import json
|
|
from pathlib import Path
|
|
from io import BytesIO
|
|
|
|
from python_repositories.adapters import MinioAdapter
|
|
|
|
from data_store.interfaces import AnnotationRepositoryInterface
|
|
from data_store.dto import (
|
|
Annotation,
|
|
DistanceEnum,
|
|
AngleEnum,
|
|
PointOfViewEnum,
|
|
ContactEnum,
|
|
)
|
|
|
|
|
|
class AnnotationRepository(MinioAdapter, AnnotationRepositoryInterface):
|
|
"""AnnotationRepository implementation using MinIO as the backend."""
|
|
|
|
encoding: str = "utf-8"
|
|
data_format: str = "json"
|
|
content_type: str = "application/json"
|
|
folder_name: str = "annotations"
|
|
|
|
@classmethod
|
|
def _object_name(cls, job_id: UUID) -> str:
|
|
"""Generate the object name for a given job ID."""
|
|
# Check input
|
|
if not isinstance(job_id, UUID):
|
|
raise ValueError("job_id must be a valid UUID.")
|
|
# Return object name
|
|
return f"{cls.folder_name}/{job_id}.{cls.data_format}"
|
|
|
|
def __enter__(self) -> AnnotationRepository:
|
|
"""Enter the runtime context related to this object."""
|
|
super().__enter__()
|
|
return self
|
|
|
|
def get(self, job_id: UUID) -> Annotation | None:
|
|
"""Retrieve an Annotation by its job ID."""
|
|
# Check input
|
|
if not isinstance(job_id, UUID):
|
|
raise ValueError("job_id must be a valid UUID.")
|
|
# Get object
|
|
object_name = self._object_name(job_id)
|
|
buffer = self._get(object_name)
|
|
if buffer is None:
|
|
return None
|
|
# Convert buffer to Annotation
|
|
try:
|
|
data_str = buffer.read().decode(self.encoding)
|
|
data = json.loads(data_str)
|
|
annotation = Annotation(
|
|
distance=DistanceEnum(data["distance"]),
|
|
angle=AngleEnum(data["angle"]),
|
|
point_of_view=PointOfViewEnum(data["point_of_view"]),
|
|
contact=ContactEnum(data["contact"]),
|
|
)
|
|
except (KeyError, ValueError) as e:
|
|
self.logger.error(
|
|
"Failed to parse Annotation from data",
|
|
object_name=object_name,
|
|
error=str(e),
|
|
)
|
|
return None
|
|
return annotation
|
|
|
|
def put(self, annotation: Annotation, job_id: UUID) -> None:
|
|
"""Update or create an Annotation."""
|
|
# Check inputs
|
|
if not isinstance(annotation, Annotation):
|
|
raise ValueError("annotation must be an Annotation instance.")
|
|
if not isinstance(job_id, UUID):
|
|
raise ValueError("job_id must be a valid UUID.")
|
|
# Convert Annotation to json
|
|
data_json = annotation.model_dump(mode="json")
|
|
data_bytes = json.dumps(data_json).encode(self.encoding)
|
|
buffer = BytesIO(data_bytes)
|
|
# Put object
|
|
object_name = self._object_name(job_id)
|
|
self._put(object_name, buffer, self.content_type)
|
|
|
|
def delete(self, job_id: UUID) -> None:
|
|
"""Delete an Annotation by its job ID."""
|
|
# Check input
|
|
if not isinstance(job_id, UUID):
|
|
raise ValueError("job_id must be a valid UUID.")
|
|
# Delete object
|
|
object_name = self._object_name(job_id)
|
|
self._delete(object_name)
|
|
|
|
def list_all(self) -> list[UUID]:
|
|
"""List all job IDs in the repository."""
|
|
# List objects in the bucket
|
|
objects = self._list_objects(prefix=self.folder_name + "/")
|
|
# Stop early if no objects
|
|
if len(objects) == 0:
|
|
return []
|
|
# Remove folder prefix and file extension
|
|
object_stems = [Path(obj).stem for obj in objects]
|
|
# Convert to list of UUIDs
|
|
job_ids = [UUID(name) for name in object_stems]
|
|
return job_ids
|