Source code for herostools.actor.archiver.zarr_archiver

# coding:utf-8

from pathlib import Path
from typing import Any

import numpy as np
import zarr
from jinja2 import Template

from .base import HERODataArchiver




class ZarrArchiver(HERODataArchiver):
    """HERODataArchiver that stores numpy array payloads in a zarr store.

    The store root and the array key within it are both derived from Jinja2
    templates rendered against the merged metadata. Opening the store in
    append mode ('a') means successive calls accumulate arrays in the same
    store directory.

    Args:
        store_template: Jinja2 template rendered to the zarr store root path.
        array_path_template: Jinja2 template rendered to the array key within the store.
        object_selector: Zenoh object selector for the devices to subscribe to.
        event_name: Name of the event.
        default_metadata: Default metadata available to both templates.
        max_retries: Number of times to retry a failed store before dropping the item.
    """

    def __init__(
        self,
        store_template: str,
        array_path_template: str,
        *args,
        **kwargs,
    ):
        self.store_template = Template(store_template)
        self.array_path_template = Template(array_path_template)
        HERODataArchiver.__init__(self, *args, **kwargs)



    def _store(self, source_name: str, payload: np.typing.NDArray[Any], metadata: dict) -> None:
        """Write the payload array into the zarr store.

        Args:
            source_name: Name of the event source (the HERO).
            payload: Numpy array to store.
            metadata: Incoming metadata merged with ``default_metadata``.
        """
        merged = self.metadata | metadata
        store_path = Path(self.store_template.render(merged))
        array_key = self.array_path_template.render(merged)
        store_path.mkdir(parents=True, exist_ok=True)
        group = zarr.open_group(str(store_path), mode="a")
        group[array_key] = payload