Source code for herostools.actor.archiver.array

# coding:utf-8

from pathlib import Path
from typing import Any

import numpy as np
from jinja2 import Template

from .base import HERODataArchiver




class ArrayArchiver(HERODataArchiver):
    """HERODataArchiver that saves numpy arrays as ``.npy`` files.

    Args:
        object_selector: Zenoh object selector for the devices to subscribe to.
        event_name: Name of the event.
        save_template: Jinja2 template used to generate the output file path from
            the payload metadata.
        split_data_array: When ``True``, each row of a 2-D array is saved as a
            separate file.  Use ``{{ _split_index }}`` in ``save_template`` to
            embed the row index in the filename.
        default_metadata: Default metadata available to the filename template.
        max_retries: Number of times to retry a failed store before dropping the item.

    Example:
        BOSS json configuration::

            {
              "_id": "my-camera-capturer",
              "classname": "herostools.archiver.ArrayArchiver",
              "arguments": {
                "object_selector": "my-camera",
                "event_name": "acquisition_data",
                "default_metadata": {
                  "file_path": "/mnt/mystorage/images"
                },
                "save_template": "{{ file_path }}/testimg-{{ '%04d' % ( frame / 2 ) |round(0, 'floor') }}-{{ frame % 2 }}.npy"
              }
            }

        Generates paths such as::

            /mnt/mystorage/images/testimg-0000-0.npy
            /mnt/mystorage/images/testimg-0000-1.npy
    """

    def __init__(
        self,
        save_template: str,
        split_data_array: bool = False,
        *args,
        **kwargs,
    ):
        self.name_template = Template(save_template)
        self.split_data_array = split_data_array
        HERODataArchiver.__init__(self, *args, **kwargs)



    def _store(self, source_name: str, payload: np.typing.NDArray[Any], metadata: dict) -> None:
        """Save the payload as a numpy ``.npy`` file.

        Args:
            source_name: Name of the event source (the HERO).
            payload: Data to save as a numpy array.
            metadata: Incoming metadata merged with ``default_metadata``.
        """
        metadata = self.metadata | metadata
        if self.split_data_array:
            for i_row, data_row in enumerate(payload):
                metadata["_split_index"] = i_row
                full_path = Path(self.name_template.render(metadata))
                full_path.parent.mkdir(parents=True, exist_ok=True)
                np.save(full_path, data_row)
        else:
            full_path = Path(self.name_template.render(metadata))
            full_path.parent.mkdir(parents=True, exist_ok=True)
            np.save(full_path, payload)