Source code for ch5mpy.attributes

from __future__ import annotations

import pickle
from typing import TYPE_CHECKING, Any

import numpy as np
from h5py._hl.attrs import AttributeManager as H5Attrs
from h5py._hl.base import KeysViewHDF5

if TYPE_CHECKING:
    pass


_SPACER = ", \n                 "
_NONE_VALUE = "__h5_NONE__"


def _is_None(v: Any) -> bool:
    if isinstance(v, np.void):
        return False

    return isinstance(v, str) and v == _NONE_VALUE


class AttributeManager:
    """Manage attributes of a File, Group or Dataset."""

    # region magic methods
    def __init__(self, attrs: H5Attrs) -> None:
        self._attrs = attrs

    def __repr__(self) -> str:
        return f"AttributeManager{{{_SPACER.join([e[0] + ': ' + str(e[1]) for e in self.as_dict().items()])}}}"

    def __getitem__(self, name: str) -> Any:
        value = self._attrs[name]

        if isinstance(value, np.void):
            return pickle.loads(value.tobytes())

        if value == _NONE_VALUE:
            return None

        return value

    def __setitem__(self, name: str, value: Any) -> Any:
        if value is None:
            self._attrs[name] = _NONE_VALUE

        else:
            try:
                self._attrs[name] = value

            except TypeError:
                # since <value> cannot be stored directly, we pickle it
                # here we use numpy's void type to allow storing bytes generated by pickle
                self._attrs[name] = np.void(pickle.dumps(value, protocol=pickle.HIGHEST_PROTOCOL))

    def __delitem__(self, name: str) -> None:
        del self._attrs[name]

    # endregion

    # region attributes
    def keys(self) -> KeysViewHDF5[str]:
        return self._attrs.keys()

    # endregion

    # region methods
    def get(self, name: str, default: Any = None) -> Any:
        return self._attrs.get(name, default)

    def set(self, **kwargs: Any) -> None:
        for k, v in kwargs.items():
            self[k] = v

    def as_dict(self) -> dict[str, Any]:
        return {k: None if _is_None(v) else v for k, v in self._attrs.items()}

    # endregion