From e55a8572e25bbcac8dce60089362362063e94a62 Mon Sep 17 00:00:00 2001 From: "Travis A. O'Brien" Date: Fri, 10 Jun 2022 12:54:56 -0400 Subject: [PATCH 1/3] `xarray.save_mfdataset` add `**kwargs` (GH6684) * adds a `**kwargs` option to `xarray.save_mfdataset` that passes through to `to_netcdf` --- doc/whats-new.rst | 4 +++ xarray/backends/api.py | 20 +++++++++-- .../tests/test_save_mfdataset_pass_kwargs.py | 34 +++++++++++++++++++ 3 files changed, 56 insertions(+), 2 deletions(-) create mode 100644 xarray/tests/test_save_mfdataset_pass_kwargs.py diff --git a/doc/whats-new.rst b/doc/whats-new.rst index da28f1dba87..833f475d38c 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -30,6 +30,10 @@ Deprecations Bug fixes ~~~~~~~~~ +- :py:meth:`xarray.save_mfdataset` now passes ``**kwargs`` on to ``to_netcdf``, + allowing the ``encoding`` and ``unlimited_dims`` options with ``save_mfdataset``. + (:issue:`6684`) + By `Travis A. O'Brien `_. Documentation ~~~~~~~~~~~~~ diff --git a/xarray/backends/api.py b/xarray/backends/api.py index d1166624569..dbd332cb9e3 100644 --- a/xarray/backends/api.py +++ b/xarray/backends/api.py @@ -1258,7 +1258,14 @@ def dump_to_store( def save_mfdataset( - datasets, paths, mode="w", format=None, groups=None, engine=None, compute=True + datasets, + paths, + mode="w", + format=None, + groups=None, + engine=None, + compute=True, + **kwargs, ): """Write multiple datasets to disk as netCDF files simultaneously. @@ -1280,6 +1287,7 @@ def save_mfdataset( these locations will be overwritten. format : {"NETCDF4", "NETCDF4_CLASSIC", "NETCDF3_64BIT", \ "NETCDF3_CLASSIC"}, optional + **kwargs : additional arguments are passed along to ``to_netcdf`` File format for the resulting netCDF file: @@ -1358,7 +1366,15 @@ def save_mfdataset( writers, stores = zip( *[ to_netcdf( - ds, path, mode, format, group, engine, compute=compute, multifile=True + ds, + path, + mode, + format, + group, + engine, + compute=compute, + multifile=True, + **kwargs, ) for ds, path, group in zip(datasets, paths, groups) ] diff --git a/xarray/tests/test_save_mfdataset_pass_kwargs.py b/xarray/tests/test_save_mfdataset_pass_kwargs.py new file mode 100644 index 00000000000..3ffbc4d877f --- /dev/null +++ b/xarray/tests/test_save_mfdataset_pass_kwargs.py @@ -0,0 +1,34 @@ +import pytest + +import xarray as xr + + +@pytest.mark.parametrize("test_type", ["baseline", "encoding_and_unlimited_dims"]) +def test_save_mfdataset_pass_kwargs(test_type): + # create a timeseries to store in a netCDF file + times = list(range(0, 1000)) + time = xr.DataArray(times, dims=("time",)) + + # create a simple dataset to write using save_mfdataset + test_ds = xr.Dataset() + test_ds["time"] = time + + # make sure the times are written as double and + # turn off fill values + encoding = dict(time=dict(dtype="double")) + unlimited_dims = ["time"] + + # set the output file name + output_path = "test.nc" + + # the test fails if save_mfdataset doesn't accept kwargs + # but it works if instead the dataset is saved using + # test_ds.to_netcdf(output_path, encoding = encoding) + if test_type == "baseline": + test_ds.to_netcdf(output_path, encoding=encoding, unlimited_dims=unlimited_dims) + elif test_type == "encoding_and_unlimited_dims": + xr.save_mfdataset( + [test_ds], [output_path], encoding=encoding, unlimited_dims=unlimited_dims + ) + else: + raise RuntimeError("Unexpected value for argument `test_type`") From 2cd2a76a3305602925cdd51058b0a295cbf7112b Mon Sep 17 00:00:00 2001 From: "Travis A. O'Brien" Date: Fri, 10 Jun 2022 13:29:25 -0400 Subject: [PATCH 2/3] Relocate test_save_mfdataset_pass_kwargs * Move test_save_mfdataset_pass_kwargs to test_backends.py * Simplify the test --- xarray/tests/test_backends.py | 23 +++++++++++++ .../tests/test_save_mfdataset_pass_kwargs.py | 34 ------------------- 2 files changed, 23 insertions(+), 34 deletions(-) delete mode 100644 xarray/tests/test_save_mfdataset_pass_kwargs.py diff --git a/xarray/tests/test_backends.py b/xarray/tests/test_backends.py index 12f6dd1508c..2c676c3ded0 100644 --- a/xarray/tests/test_backends.py +++ b/xarray/tests/test_backends.py @@ -3706,6 +3706,29 @@ def test_save_mfdataset_pathlib_roundtrip(self) -> None: ) as actual: assert_identical(actual, original) + def test_save_mfdataset_pass_kwargs(): + # create a timeseries to store in a netCDF file + times = [0, 1] + time = xr.DataArray(times, dims=("time",)) + + # create a simple dataset to write using save_mfdataset + test_ds = xr.Dataset() + test_ds["time"] = time + + # make sure the times are written as double and + # turn off fill values + encoding = dict(time=dict(dtype="double")) + unlimited_dims = ["time"] + + # set the output file name + output_path = "test.nc" + + # attempt to write the dataset with the encoding and unlimited args + # passed through + xr.save_mfdataset( + [test_ds], [output_path], encoding=encoding, unlimited_dims=unlimited_dims + ) + def test_open_and_do_math(self) -> None: original = Dataset({"foo": ("x", np.random.randn(10))}) with create_tmp_file() as tmp: diff --git a/xarray/tests/test_save_mfdataset_pass_kwargs.py b/xarray/tests/test_save_mfdataset_pass_kwargs.py deleted file mode 100644 index 3ffbc4d877f..00000000000 --- a/xarray/tests/test_save_mfdataset_pass_kwargs.py +++ /dev/null @@ -1,34 +0,0 @@ -import pytest - -import xarray as xr - - -@pytest.mark.parametrize("test_type", ["baseline", "encoding_and_unlimited_dims"]) -def test_save_mfdataset_pass_kwargs(test_type): - # create a timeseries to store in a netCDF file - times = list(range(0, 1000)) - time = xr.DataArray(times, dims=("time",)) - - # create a simple dataset to write using save_mfdataset - test_ds = xr.Dataset() - test_ds["time"] = time - - # make sure the times are written as double and - # turn off fill values - encoding = dict(time=dict(dtype="double")) - unlimited_dims = ["time"] - - # set the output file name - output_path = "test.nc" - - # the test fails if save_mfdataset doesn't accept kwargs - # but it works if instead the dataset is saved using - # test_ds.to_netcdf(output_path, encoding = encoding) - if test_type == "baseline": - test_ds.to_netcdf(output_path, encoding=encoding, unlimited_dims=unlimited_dims) - elif test_type == "encoding_and_unlimited_dims": - xr.save_mfdataset( - [test_ds], [output_path], encoding=encoding, unlimited_dims=unlimited_dims - ) - else: - raise RuntimeError("Unexpected value for argument `test_type`") From f4cb5a7b05e2604ec153ece9f45c2d03b4174fea Mon Sep 17 00:00:00 2001 From: "Travis A. O'Brien" Date: Fri, 10 Jun 2022 13:52:02 -0400 Subject: [PATCH 3/3] Update xarray/tests/test_backends.py Co-authored-by: Anderson Banihirwe --- xarray/tests/test_backends.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/xarray/tests/test_backends.py b/xarray/tests/test_backends.py index 2c676c3ded0..bbfae73eadd 100644 --- a/xarray/tests/test_backends.py +++ b/xarray/tests/test_backends.py @@ -3706,7 +3706,7 @@ def test_save_mfdataset_pathlib_roundtrip(self) -> None: ) as actual: assert_identical(actual, original) - def test_save_mfdataset_pass_kwargs(): + def test_save_mfdataset_pass_kwargs(self): # create a timeseries to store in a netCDF file times = [0, 1] time = xr.DataArray(times, dims=("time",))