Source code for herostools.actor.archiver.zarr_archiver

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. Use 

        Args:
            source_name: Name of the event source (the HERO).
            payload: Numpy array to store.
            metadata: Incoming metadata merged with ``default_metadata``. 
                Dicts provided with ``group-attributes`` or ``array-attributes`` will be merged in the 
                respective Zarr attributes. Skipped if any Zarr attribute is invalid.
        """
        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
        # Try Merge attributes 
        try: 
            group.update_attributes(merged.get('group-attributes', {}))
            group[array_key].update_attributes(merged.get('array-attributes', {}))
        except ValueError:
            pass