"""HDF5 storage backend for pytcl data persistence.
Provides efficient storage of large numerical arrays using HDF5.
h5py is a core pytcl dependency (see pyproject.toml); no optional
install step is needed.
"""
import json
from pathlib import Path
from types import TracebackType
from typing import Any, Dict, List, Optional, Union
import h5py
import numpy as np
from numpy.typing import ArrayLike, NDArray
from pytcl.io.storage import StorageBackend
[docs]
class HDF5Storage(StorageBackend):
"""HDF5-based storage backend.
Efficiently stores large arrays and structured data using HDF5 format.
Ideal for numerical data, model coefficients, and time series.
Examples
--------
>>> from pytcl.io import HDF5Storage
>>> with HDF5Storage() as store: # doctest: +SKIP
... store.open("data.h5", mode="w")
... store.store_array("gravity/egm96", coefficients)
... store.store_scalar("header/version", 1)
... result = store.retrieve_array("gravity/egm96")
"""
[docs]
def __init__(self) -> None:
self._file = None
self._path = None
self._mode = None
[docs]
def open(self, path: str, mode: str = "r") -> None:
"""Open an HDF5 file.
Parameters
----------
path : str
Path to HDF5 file
mode : str, optional
'r' (read), 'w' (write), 'a' (append). Default is 'r'.
"""
self._path = Path(path)
self._mode = mode
self._file = h5py.File(str(self._path), mode=mode)
[docs]
def close(self) -> None:
"""Close the HDF5 file."""
if self._file is not None:
self._file.close()
self._file = None
def __enter__(self) -> "HDF5Storage":
return self
def __exit__(
self,
exc_type: type[BaseException] | None,
exc_val: BaseException | None,
exc_tb: TracebackType | None,
) -> None:
self.close()
[docs]
def store_array(
self,
name: str,
data: ArrayLike,
metadata: Optional[Dict[str, Any]] = None,
) -> None:
"""Store a numpy array as an HDF5 dataset.
Parameters
----------
name : str
Dataset path (e.g., "groups/subgroup/data")
data : ArrayLike
Array to store
metadata : dict, optional
Metadata stored as HDF5 attributes
Notes
-----
Storing over an existing name replaces it, which is the contract
``StorageBackend`` defines and what ``SQLStorage`` already did. This
used to let h5py raise ``ValueError`` instead, so the same code worked
against one backend and failed against the other (gh-21).
"""
if self._file is None:
raise RuntimeError("Storage file not open. Call open() first.")
# Ensure parent groups exist
self._ensure_groups(name)
# Replace rather than let h5py refuse an existing name. Deleting the
# dataset drops its attributes too, so the metadata below is a
# replacement rather than a merge -- as the base class specifies.
if name in self._file:
del self._file[name]
# Store array
arr = np.asarray(data)
dataset = self._file.create_dataset(name, data=arr)
# Store metadata as attributes
if metadata:
for key, value in metadata.items():
try:
dataset.attrs[key] = value
except (TypeError, ValueError):
# For non-serializable objects, store as JSON string
dataset.attrs[key] = json.dumps(str(value))
[docs]
def retrieve_array(self, name: str) -> NDArray[Any]:
"""Retrieve a stored array.
Parameters
----------
name : str
Dataset path
Returns
-------
ndarray
The stored array
"""
if self._file is None:
raise RuntimeError("Storage file not open. Call open() first.")
if name not in self._file:
raise KeyError(f"Dataset '{name}' not found in HDF5 file")
return np.array(self._file[name])
[docs]
def store_scalar(
self,
name: str,
value: Union[int, float, str, bool],
metadata: Optional[Dict[str, Any]] = None,
) -> None:
"""Store a scalar value, replacing any scalar already under that name.
Parameters
----------
name : str
Scalar name/path
value : scalar
Value to store
metadata : dict, optional
Associated metadata
"""
if self._file is None:
raise RuntimeError("Storage file not open. Call open() first.")
# Ensure parent groups exist
self._ensure_groups(name)
# Replace on collision, matching store_array and SQLStorage (the
# gh-21 contract; this method was missed by that fix).
if name in self._file:
del self._file[name]
# Store as dataset with shape ()
dataset = self._file.create_dataset(name, data=value)
if metadata:
for key, val in metadata.items():
try:
dataset.attrs[key] = val
except (TypeError, ValueError):
dataset.attrs[key] = json.dumps(str(val))
[docs]
def retrieve_scalar(self, name: str) -> Union[int, float, str, bool]:
"""Retrieve a scalar value.
Parameters
----------
name : str
Scalar name/path
Returns
-------
Scalar value
"""
if self._file is None:
raise RuntimeError("Storage file not open. Call open() first.")
if name not in self._file:
raise KeyError(f"Scalar '{name}' not found in HDF5 file")
value = self._file[name][()]
# Handle bytes (HDF5 stores strings as bytes)
if isinstance(value, bytes):
return value.decode("utf-8")
return value.item() if hasattr(value, "item") else value
[docs]
def store_group(self, name: str, metadata: Optional[Dict[str, Any]] = None) -> None:
"""Create a group for organizing related datasets.
Parameters
----------
name : str
Group path
metadata : dict, optional
Group-level metadata
"""
if self._file is None:
raise RuntimeError("Storage file not open. Call open() first.")
# Create groups recursively if they don't already exist
if name not in self._file:
self._file.create_group(name)
if metadata:
group = self._file[name]
for key, val in metadata.items():
try:
group.attrs[key] = val
except (TypeError, ValueError):
group.attrs[key] = json.dumps(str(val))
[docs]
def list_keys(self, group: str = "/") -> List[str]:
"""List datasets and groups in a location.
Parameters
----------
group : str, optional
Group path. Default is root.
Returns
-------
list of str
Keys in the group
"""
if self._file is None:
raise RuntimeError("Storage file not open. Call open() first.")
try:
return list(self._file[group].keys())
except KeyError:
raise KeyError(f"Group '{group}' not found in HDF5 file")
[docs]
def delete(self, name: str) -> None:
"""Delete a dataset or group.
Parameters
----------
name : str
Dataset/group path
"""
if self._file is None:
raise RuntimeError("Storage file not open. Call open() first.")
if name in self._file:
del self._file[name]
[docs]
def flush(self) -> None:
"""Ensure all data is written to disk."""
if self._file is not None:
self._file.flush()
def _ensure_groups(self, path: str, is_group: bool = False) -> None:
"""Ensure parent groups exist, creating them if needed.
Parameters
----------
path : str
Full path to dataset or group
is_group : bool, optional
If True, treat path as group. Default treats as dataset path.
"""
if "/" not in path:
return
# Get parent path
if is_group:
parent = path
else:
parent = path.rsplit("/", 1)[0]
# Both callers reject a closed file before reaching here; this states
# the precondition for the type checker rather than re-validating it.
assert self._file is not None
if parent and parent != "/" and parent not in self._file:
self._file.create_group(parent)