"""Collection of helper functions related to NWB."""
import importlib
import uuid
import warnings
from contextlib import contextmanager
from copy import deepcopy
from datetime import datetime
from pathlib import Path
from typing import Literal
from pydantic import FilePath
from pynwb import NWBFile, read_nwb
from pynwb.file import Subject
from . import (
BACKEND_NWB_IO,
BackendConfiguration,
configure_backend,
get_default_backend_configuration,
)
from ...utils.dict import DeepDict, load_dict_from_file
from ...utils.json_schema import validate_metadata
[docs]
def get_module(nwbfile: NWBFile, name: str, description: str = None):
"""
Check if processing module exists. If not, create it. Then return module.
Parameters
----------
nwbfile : NWBFile
The NWB file to check or add the module to.
name : str
The name of the processing module.
description : str, optional
Description of the module. Only used if creating a new module.
Returns
-------
ProcessingModule
The existing or newly created processing module.
"""
if name in nwbfile.processing:
if description is not None and nwbfile.processing[name].description != description:
existing_description = nwbfile.processing[name].description
warnings.warn(
f"Processing module '{name}' already exists with a different description. "
f"The new description will be ignored.\n"
f" Existing: '{existing_description}'\n"
f" Provided: '{description}'\n"
f"To fix this, ensure all calls to get_module() for '{name}' use the same description, "
f"or omit the description parameter to use the existing one."
)
return nwbfile.processing[name]
else:
if description is None:
description = "No description."
return nwbfile.create_processing_module(name=name, description=description)
def _add_device_to_nwbfile(
*,
nwbfile: NWBFile,
device_metadata: dict,
):
"""
Add a device to an NWBFile.
If a device with the same name already exists, the existing device is
returned without creating a duplicate.
Parameters
----------
nwbfile : NWBFile
The NWB file to add the device to.
device_metadata : dict
Dictionary describing the device. Must contain at least a ``"name"`` key.
Returns
-------
Device
The Device object (either newly created or existing).
"""
device_name = device_metadata["name"]
if device_name in nwbfile.devices:
return nwbfile.devices[device_name]
device = nwbfile.create_device(**device_metadata)
return device
def _attempt_cleanup_of_existing_nwbfile(nwbfile_path: Path) -> None:
if not nwbfile_path.exists():
return
try:
nwbfile_path.unlink()
# Windows in particular can encounter errors at this step
except PermissionError: # pragma: no cover
message = f"Unable to remove NWB file located at {nwbfile_path.absolute()}! Please remove it manually."
warnings.warn(message=message, stacklevel=2)
[docs]
@contextmanager
def make_or_load_nwbfile(
nwbfile_path: FilePath | None = None,
nwbfile: NWBFile | None = None,
metadata: dict | None = None,
overwrite: bool = False,
backend: Literal["hdf5", "zarr"] = "hdf5",
verbose: bool = False,
):
"""
Context for automatically handling decision of write vs. append for writing an NWBFile.
Parameters
----------
nwbfile_path: FilePath
Path for where to write or load (if overwrite=False) the NWBFile.
If specified, the context will always write to this location.
nwbfile: NWBFile, optional
An in-memory NWBFile object to write to the location.
metadata: dict, optional
Metadata dictionary with information used to create the NWBFile when one does not exist or overwrite=True.
overwrite: bool, default: False
Whether to overwrite the NWBFile if one exists at the nwbfile_path.
The default is False (append mode).
backend : "hdf5" or "zarr", default: "hdf5"
The type of backend used to create the file.
verbose: bool, default: True
If 'nwbfile_path' is specified, informs user after a successful write operation.
"""
from . import BACKEND_NWB_IO
nwbfile_path_is_provided = nwbfile_path is not None
nwbfile_path_in = Path(nwbfile_path) if nwbfile_path_is_provided else None
nwbfile_is_provided = nwbfile is not None
nwbfile_in = nwbfile if nwbfile_is_provided else None
backend_io_class = BACKEND_NWB_IO[backend]
assert not (nwbfile_path is None and nwbfile is None and metadata is None), (
"You must specify either an 'nwbfile_path', or an in-memory 'nwbfile' object, "
"or provide the metadata for creating one."
)
assert not (overwrite is False and nwbfile_path_in and nwbfile_path_in.exists() and nwbfile is not None), (
"'nwbfile_path' exists at location, 'overwrite' is False (append mode), but an in-memory 'nwbfile' object was "
"passed! Cannot reconcile which nwbfile object to write."
)
if overwrite is False and backend == "zarr":
# TODO: remove when https://github.com/hdmf-dev/hdmf-zarr/issues/182 is resolved
raise NotImplementedError("Appending a Zarr file is not yet supported!")
load_kwargs = dict()
file_initially_exists = nwbfile_path_in.exists() if nwbfile_path_is_provided else False
append_mode = file_initially_exists and not overwrite
if nwbfile_path_is_provided:
load_kwargs.update(path=str(nwbfile_path_in))
if append_mode:
load_kwargs.update(mode="r+", load_namespaces=True)
# Check if the selected backend is the backend of the file in nwfile_path
backends_that_can_read = [
backend_name
for backend_name, backend_io_class in BACKEND_NWB_IO.items()
if backend_io_class.can_read(path=str(nwbfile_path_in))
]
# Future-proofing: raise an error if more than one backend can read the file
assert (
len(backends_that_can_read) <= 1
), "More than one backend is capable of reading the file! Please raise an issue describing your file."
if backend not in backends_that_can_read:
raise IOError(
f"The chosen backend ('{backend}') is unable to read the file! "
f"Please select '{backends_that_can_read[0]}' instead."
)
else:
load_kwargs.update(mode="w")
io = backend_io_class(**load_kwargs)
read_nwbfile = nwbfile_path_is_provided and append_mode
create_nwbfile = not read_nwbfile and not nwbfile_is_provided
nwbfile_loaded_succesfully = True
nwbfile_written_succesfully = True
try:
if nwbfile_is_provided:
nwbfile = nwbfile_in
elif read_nwbfile:
nwbfile = io.read()
elif create_nwbfile:
if metadata is None:
error_msg = "Metadata is required for creating an nwbfile "
raise ValueError(error_msg)
default_metadata = get_default_nwbfile_metadata()
default_metadata.deep_update(metadata)
nwbfile = make_nwbfile_from_metadata(metadata=metadata)
yield nwbfile
except Exception as load_error:
nwbfile_loaded_succesfully = False
raise load_error
finally:
if nwbfile_path_is_provided and nwbfile_loaded_succesfully:
try:
io.write(nwbfile)
if verbose:
print(f"NWB file saved at {nwbfile_path_in}!")
except Exception as write_error:
nwbfile_written_succesfully = False
raise write_error
finally:
io.close()
del io
if not nwbfile_written_succesfully:
_attempt_cleanup_of_existing_nwbfile(nwbfile_path=nwbfile_path_in)
elif nwbfile_path_is_provided and not nwbfile_loaded_succesfully:
# The instantiation of the IO object can itself create a file
_attempt_cleanup_of_existing_nwbfile(nwbfile_path=nwbfile_path_in)
else:
# This is the case where nwbfile is provided but not nwbfile_path
# Note that io never gets created in this case, so no need to close or delete it
pass
# Final attempt to cleanup an unintended file creation, just to be sure
any_load_or_write_error = not nwbfile_loaded_succesfully or not nwbfile_written_succesfully
file_was_freshly_created = not file_initially_exists and nwbfile_path_is_provided and nwbfile_path_in.exists()
attempt_to_cleanup = any_load_or_write_error and file_was_freshly_created
if attempt_to_cleanup:
_attempt_cleanup_of_existing_nwbfile(nwbfile_path=nwbfile_path_in)
def _resolve_backend(
backend: Literal["hdf5", "zarr"] | None = None,
backend_configuration: BackendConfiguration | None = None,
) -> Literal["hdf5"]:
"""
Resolve the backend to use for writing the NWBFile.
Parameters
----------
backend: {"hdf5", "zarr"}, optional
backend_configuration: BackendConfiguration, optional
Returns
-------
backend: {"hdf5", "zarr"}
"""
if backend is not None and backend_configuration is not None:
if backend == backend_configuration.backend:
warnings.warn(
f"Both `backend` and `backend_configuration` were specified as type '{backend}'. "
"To suppress this warning, specify only `backend_configuration`."
)
else:
raise ValueError(
f"Both `backend` and `backend_configuration` were specified and are conflicting."
f"{backend=}, {backend_configuration.backend=}."
"These values must match. To suppress this error, specify only `backend_configuration`."
)
if backend is None:
backend = backend_configuration.backend if backend_configuration is not None else "hdf5"
return backend
[docs]
def repack_nwbfile(
*,
nwbfile_path: Path,
export_nwbfile_path: Path,
export_backend: Literal["hdf5", "zarr", None] = None,
):
"""
Repack an NWBFile with a new backend configuration.
Parameters
----------
nwbfile_path : Path
Path to the NWB file to be repacked.
export_nwbfile_path : Path
Path to export the repacked NWB file.
export_backend : {"hdf5", "zarr", None}, default: None
The type of backend used to write the repacked file. If None, the same backend as the input file is used.
"""
# Read the file using read_nwb (automatically detects backend)
nwbfile = read_nwb(nwbfile_path)
# If no export backend specified, detect from the read_io attribute
if export_backend is None:
# Determine the backend from the IO class that was used to read the file
io_class_name = nwbfile.read_io.__class__.__name__
# Map IO class names to backend names
if "NWBHDF5IO" in io_class_name:
export_backend = "hdf5"
elif "NWBZarrIO" in io_class_name:
export_backend = "zarr"
else:
raise ValueError(f"Unable to determine backend from IO class: {io_class_name}")
backend_configuration = get_default_backend_configuration(nwbfile=nwbfile, backend=export_backend)
configure_and_write_nwbfile(
nwbfile=nwbfile,
backend_configuration=backend_configuration,
nwbfile_path=export_nwbfile_path,
backend=export_backend,
)