def test_write_dataset_error(tmp_path: Path):
"""Raise ValueError if method isn't implemented in the DataIo class."""
file_path = tmp_path / "dummy.foo"
with pytest.raises(ValueError,
match="Cannot save data with format: 'foo'"):
save_dataset(str(file_path), "", "foo") # type:ignore
file_path.touch()
with pytest.raises(ValueError,
match="Cannot read data with format: 'foo'"):
load_dataset(str(file_path))
def test_load_dataset(tmp_path: Path):
"""All args and kwargs are passes correctly."""
file_path = tmp_path / "dummy.mock"
file_path.write_text("mock")
result = load_dataset(str(file_path), dummy_arg="baz")
assert result == {"file_name": str(file_path), "dummy_arg": "baz"}
def test_load_dataset(tmp_path: Path):
"""All args and kwargs are passes correctly."""
file_path = tmp_path / "dummy.mock"
file_path.write_text("mock")
result: dict[str, Any] = {}
dataset = load_dataset(file_path, result_container=result, dummy_arg="baz")
assert result == {"file_name": file_path.as_posix(), "dummy_arg": "baz"}
assert np.all(dataset.data == xr.DataArray([1, 2]).to_dataset(
name="data").data)
assert dataset.source_path == file_path.as_posix()
def _load_datasets(dataset_mappable: DatasetMappable,
index: int = 1) -> dict[str, xr.Dataset]:
"""Implement functionality for ``load_datasets`` and internal use.
Parameters
----------
dataset_mappable : DatasetMappable
Instance of ``DatasetMappable`` that can be used to create a dataset mapping.
index : int
Index used to create key and ``source_path`` if not present.
, by default 1
Returns
-------
dict[str, xr.Dataset]
Mapping of datasets to initialize :class:`DatasetMapping`.
Raises
------
TypeError
If the type of ``dataset_mappable`` is not explicitly supported.
"""
dataset_mapping = {}
if isinstance(dataset_mappable, (str, Path)):
dataset_mapping[Path(dataset_mappable).stem] = load_dataset(
dataset_mappable)
elif isinstance(dataset_mappable, (xr.Dataset, xr.DataArray)):
if isinstance(dataset_mappable, xr.DataArray):
dataset_mappable: xr.Dataset = dataset_mappable.to_dataset( # type:ignore[no-redef]
name="data")
if "source_path" not in dataset_mappable.attrs:
dataset_mappable.attrs["source_path"] = f"dataset_{index}.nc"
dataset_mapping[Path(
dataset_mappable.source_path).stem] = dataset_mappable
elif isinstance(dataset_mappable, Sequence):
for index, dataset in enumerate(dataset_mappable, start=1):
key, value = next(
iter(_load_datasets(dataset, index=index).items()))
dataset_mapping[key] = value
elif isinstance(dataset_mappable, Mapping):
for key, dataset in dataset_mappable.items():
_, value = next(iter(_load_datasets(dataset).items()))
dataset_mapping[key] = value
else:
raise TypeError(
f"Type '{type(dataset_mappable).__name__}' for 'dataset_mappable' of value "
f"'{dataset_mappable}' is not supported."
f"\nSupported types are:\n {DatasetMappable}.")
return dataset_mapping