From 914a2210b40e8ba9dcacda9cb46cdc01b956c244 Mon Sep 17 00:00:00 2001 From: Joseph Hamman Date: Sat, 2 Feb 2019 15:49:18 -0800 Subject: [PATCH 1/5] add h5netcdf+dask tests --- xarray/tests/test_backends.py | 39 +++++++++++++++++++++++++++++++++++ 1 file changed, 39 insertions(+) diff --git a/xarray/tests/test_backends.py b/xarray/tests/test_backends.py index 55e4eb7c8db..31881070c5f 100644 --- a/xarray/tests/test_backends.py +++ b/xarray/tests/test_backends.py @@ -1916,6 +1916,45 @@ def test_dump_encodings_h5py(self): assert actual.x.encoding['compression_opts'] is None +@requires_h5netcdf +@requires_dask +@pytest.mark.filterwarnings('ignore:deallocating CachingFileManager') +class TestH5NetCDFViaDaskData(TestH5NetCDFData): + + @contextlib.contextmanager + def roundtrip(self, data, save_kwargs={}, open_kwargs={}, + allow_cleanup_failure=False): + with TestH5NetCDFData.roundtrip( + self, data, save_kwargs, open_kwargs, + allow_cleanup_failure) as ds: + yield ds.chunk() + + def test_dataset_caching(self): + # caching behavior differs for dask + pass + + @pytest.mark.xfail(reason="Failing to round trip unlimited dims") + def test_encoding_unlimited_dims(self): + # TODO: this should pass + super(TestH5NetCDFViaDaskData, self).test_encoding_unlimited_dims() + + def test_write_inconsistent_chunks(self): + # Construct two variables with the same dimensions, but different + # chunk sizes. + x = da.zeros((100, 100), dtype='f4', chunks=(50, 100)) + x = DataArray(data=x, dims=('lat', 'lon'), name='x') + x.encoding['chunksizes'] = (50, 100) + x.encoding['original_shape'] = (100, 100) + y = da.ones((100, 100), dtype='f4', chunks=(100, 50)) + y = DataArray(data=y, dims=('lat', 'lon'), name='y') + y.encoding['chunksizes'] = (100, 50) + y.encoding['original_shape'] = (100, 100) + # Put them both into the same dataset + ds = Dataset({'x': x, 'y': y}) + with self.roundtrip(ds) as actual: + assert actual['x'].encoding['chunksizes'] == (50, 100) + assert actual['y'].encoding['chunksizes'] == (100, 50) + @pytest.fixture(params=['scipy', 'netcdf4', 'h5netcdf', 'pynio']) def readengine(request): return request.param From 98850eee8245b66b29be79cd1b2e81a2758fe9a8 Mon Sep 17 00:00:00 2001 From: Joseph Hamman Date: Sat, 2 Feb 2019 15:50:56 -0800 Subject: [PATCH 2/5] pep8 --- xarray/tests/test_backends.py | 1 + 1 file changed, 1 insertion(+) diff --git a/xarray/tests/test_backends.py b/xarray/tests/test_backends.py index 31881070c5f..2665691dfda 100644 --- a/xarray/tests/test_backends.py +++ b/xarray/tests/test_backends.py @@ -1955,6 +1955,7 @@ def test_write_inconsistent_chunks(self): assert actual['x'].encoding['chunksizes'] == (50, 100) assert actual['y'].encoding['chunksizes'] == (100, 50) + @pytest.fixture(params=['scipy', 'netcdf4', 'h5netcdf', 'pynio']) def readengine(request): return request.param From 30ae6c7a7f502b209d1eedacd224b1c68ec131ce Mon Sep 17 00:00:00 2001 From: Joseph Hamman Date: Sun, 3 Feb 2019 20:56:48 -0800 Subject: [PATCH 3/5] pass encoding through to _replace_vars_and_dims in ds.chunk() --- xarray/core/dataset.py | 15 ++++++++++----- xarray/tests/test_backends.py | 23 +++++++++++++++-------- 2 files changed, 25 insertions(+), 13 deletions(-) diff --git a/xarray/core/dataset.py b/xarray/core/dataset.py index 8863dedb7db..16290f92495 100644 --- a/xarray/core/dataset.py +++ b/xarray/core/dataset.py @@ -675,8 +675,8 @@ def _from_vars_and_coord_names(cls, variables, coord_names, attrs=None): return cls._construct_direct(variables, coord_names, dims, attrs) def _replace_vars_and_dims(self, variables, coord_names=None, dims=None, - attrs=__default_attrs, indexes=None, - inplace=False): + attrs=__default_attrs, encoding=None, + indexes=None, inplace=False): """Fastpath constructor for internal use. Preserves coord names and attributes. If not provided explicitly, @@ -691,6 +691,7 @@ def _replace_vars_and_dims(self, variables, coord_names=None, dims=None, variables : OrderedDict coord_names : set or None, optional attrs : OrderedDict or None, optional + encoding : OrderedDict or None, optional Returns ------- @@ -705,6 +706,8 @@ def _replace_vars_and_dims(self, variables, coord_names=None, dims=None, self._coord_names = coord_names if attrs is not self.__default_attrs: self._attrs = attrs + if encoding is not None: + self._encoding = encoding self._indexes = indexes obj = self else: @@ -712,8 +715,9 @@ def _replace_vars_and_dims(self, variables, coord_names=None, dims=None, coord_names = self._coord_names.copy() if attrs is self.__default_attrs: attrs = self._attrs_copy() - obj = self._construct_direct( - variables, coord_names, dims, attrs, indexes) + obj = self._construct_direct(variables, coord_names=coord_names, + dims=dims, attrs=attrs, + indexes=indexes, encoding=encoding) return obj def _replace_indexes(self, indexes): @@ -1410,7 +1414,8 @@ def maybe_chunk(name, var, chunks): variables = OrderedDict([(k, maybe_chunk(k, v, chunks)) for k, v in self.variables.items()]) - return self._replace_vars_and_dims(variables) + return self._replace_vars_and_dims(variables, attrs=self.attrs, + encoding=self.encoding) def _validate_indexers(self, indexers): """ Here we make sure diff --git a/xarray/tests/test_backends.py b/xarray/tests/test_backends.py index 2665691dfda..410dfbb1a16 100644 --- a/xarray/tests/test_backends.py +++ b/xarray/tests/test_backends.py @@ -1194,6 +1194,17 @@ def test_read_variable_len_strings(self): with open_dataset(tmp_file, **kwargs) as actual: assert_identical(expected, actual) + def test_encoding_unlimited_dims(self): + ds = Dataset({'x': ('y', np.arange(10.0))}) + with self.roundtrip(ds, + save_kwargs=dict(unlimited_dims=['y'])) as actual: + assert actual.encoding['unlimited_dims'] == set('y') + assert_equal(ds, actual) + ds.encoding = {'unlimited_dims': ['y']} + with self.roundtrip(ds) as actual: + assert actual.encoding['unlimited_dims'] == set('y') + assert_equal(ds, actual) + @requires_netCDF4 class TestNetCDF4Data(NetCDF4Base): @@ -1278,10 +1289,11 @@ class TestNetCDF4ViaDaskData(TestNetCDF4Data): @contextlib.contextmanager def roundtrip(self, data, save_kwargs={}, open_kwargs={}, allow_cleanup_failure=False): + open_kwargs['chunks'] = open_kwargs.get('chunks', 1) with TestNetCDF4Data.roundtrip( self, data, save_kwargs, open_kwargs, allow_cleanup_failure) as ds: - yield ds.chunk() + yield ds def test_unsorted_index_raises(self): # Skip when using dask because dask rewrites indexers to getitem, @@ -1309,7 +1321,6 @@ def test_write_inconsistent_chunks(self): assert actual['x'].encoding['chunksizes'] == (50, 100) assert actual['y'].encoding['chunksizes'] == (100, 50) - @requires_zarr class ZarrBase(CFEncodedBase): @@ -1924,20 +1935,16 @@ class TestH5NetCDFViaDaskData(TestH5NetCDFData): @contextlib.contextmanager def roundtrip(self, data, save_kwargs={}, open_kwargs={}, allow_cleanup_failure=False): + open_kwargs['chunks'] = open_kwargs.get('chunks', 1) with TestH5NetCDFData.roundtrip( self, data, save_kwargs, open_kwargs, allow_cleanup_failure) as ds: - yield ds.chunk() + yield ds def test_dataset_caching(self): # caching behavior differs for dask pass - @pytest.mark.xfail(reason="Failing to round trip unlimited dims") - def test_encoding_unlimited_dims(self): - # TODO: this should pass - super(TestH5NetCDFViaDaskData, self).test_encoding_unlimited_dims() - def test_write_inconsistent_chunks(self): # Construct two variables with the same dimensions, but different # chunk sizes. From 25549ac40f46b7636cf1d243371d3398979738b7 Mon Sep 17 00:00:00 2001 From: Joseph Hamman Date: Sun, 3 Feb 2019 21:29:41 -0800 Subject: [PATCH 4/5] lint --- xarray/tests/test_backends.py | 1 + 1 file changed, 1 insertion(+) diff --git a/xarray/tests/test_backends.py b/xarray/tests/test_backends.py index 410dfbb1a16..1d19eae234e 100644 --- a/xarray/tests/test_backends.py +++ b/xarray/tests/test_backends.py @@ -1321,6 +1321,7 @@ def test_write_inconsistent_chunks(self): assert actual['x'].encoding['chunksizes'] == (50, 100) assert actual['y'].encoding['chunksizes'] == (100, 50) + @requires_zarr class ZarrBase(CFEncodedBase): From 538dd45e7304e1e64406f52365b96c2abe706b76 Mon Sep 17 00:00:00 2001 From: Joseph Hamman Date: Wed, 6 Feb 2019 17:18:39 -0800 Subject: [PATCH 5/5] _kwargs=None in roundtrip methods --- xarray/tests/test_backends.py | 56 +++++++++++++++++++++++++++-------- 1 file changed, 44 insertions(+), 12 deletions(-) diff --git a/xarray/tests/test_backends.py b/xarray/tests/test_backends.py index 89cbfa0d1c7..d0ca23c2792 100644 --- a/xarray/tests/test_backends.py +++ b/xarray/tests/test_backends.py @@ -171,8 +171,12 @@ def create_store(self): raise NotImplementedError @contextlib.contextmanager - def roundtrip(self, data, save_kwargs={}, open_kwargs={}, + def roundtrip(self, data, save_kwargs=None, open_kwargs=None, allow_cleanup_failure=False): + if save_kwargs is None: + save_kwargs = {} + if open_kwargs is None: + open_kwargs = {} with create_tmp_file( allow_cleanup_failure=allow_cleanup_failure) as path: self.save(data, path, **save_kwargs) @@ -180,8 +184,12 @@ def roundtrip(self, data, save_kwargs={}, open_kwargs={}, yield ds @contextlib.contextmanager - def roundtrip_append(self, data, save_kwargs={}, open_kwargs={}, + def roundtrip_append(self, data, save_kwargs=None, open_kwargs=None, allow_cleanup_failure=False): + if save_kwargs is None: + save_kwargs = {} + if open_kwargs is None: + open_kwargs = {} with create_tmp_file( allow_cleanup_failure=allow_cleanup_failure) as path: for i, key in enumerate(data.variables): @@ -1287,9 +1295,13 @@ def test_autoclose_future_warning(self): @pytest.mark.filterwarnings('ignore:deallocating CachingFileManager') class TestNetCDF4ViaDaskData(TestNetCDF4Data): @contextlib.contextmanager - def roundtrip(self, data, save_kwargs={}, open_kwargs={}, + def roundtrip(self, data, save_kwargs=None, open_kwargs=None, allow_cleanup_failure=False): - open_kwargs['chunks'] = open_kwargs.get('chunks', 1) + if open_kwargs is None: + open_kwargs = {} + if save_kwargs is None: + save_kwargs = {} + open_kwargs.setdefault('chunks', -1) with TestNetCDF4Data.roundtrip( self, data, save_kwargs, open_kwargs, allow_cleanup_failure) as ds: @@ -1341,15 +1353,19 @@ def open(self, store_target, **kwargs): yield ds @contextlib.contextmanager - def roundtrip(self, data, save_kwargs={}, open_kwargs={}, + def roundtrip(self, data, save_kwargs=None, open_kwargs=None, allow_cleanup_failure=False): + if save_kwargs is None: + save_kwargs = {} + if open_kwargs is None: + open_kwargs = {} with self.create_zarr_target() as store_target: self.save(data, store_target, **save_kwargs) with self.open(store_target, **open_kwargs) as ds: yield ds @contextlib.contextmanager - def roundtrip_append(self, data, save_kwargs={}, open_kwargs={}, + def roundtrip_append(self, data, save_kwargs=None, open_kwargs=None, allow_cleanup_failure=False): pytest.skip("zarr backend does not support appending") @@ -1630,8 +1646,12 @@ def create_store(self): yield backends.ScipyDataStore(fobj, 'w') @contextlib.contextmanager - def roundtrip(self, data, save_kwargs={}, open_kwargs={}, + def roundtrip(self, data, save_kwargs=None, open_kwargs=None, allow_cleanup_failure=False): + if save_kwargs is None: + save_kwargs = {} + if open_kwargs is None: + open_kwargs = {} with create_tmp_file() as tmp_file: with open(tmp_file, 'wb') as f: self.save(data, f, **save_kwargs) @@ -1932,9 +1952,13 @@ def test_dump_encodings_h5py(self): class TestH5NetCDFViaDaskData(TestH5NetCDFData): @contextlib.contextmanager - def roundtrip(self, data, save_kwargs={}, open_kwargs={}, + def roundtrip(self, data, save_kwargs=None, open_kwargs=None, allow_cleanup_failure=False): - open_kwargs['chunks'] = open_kwargs.get('chunks', 1) + if save_kwargs is None: + save_kwargs = {} + if open_kwargs is None: + open_kwargs = {} + open_kwargs.setdefault('chunks', -1) with TestH5NetCDFData.roundtrip( self, data, save_kwargs, open_kwargs, allow_cleanup_failure) as ds: @@ -2146,7 +2170,7 @@ def create_store(self): yield Dataset() @contextlib.contextmanager - def roundtrip(self, data, save_kwargs={}, open_kwargs={}, + def roundtrip(self, data, save_kwargs=None, open_kwargs=None, allow_cleanup_failure=False): yield data.chunk() @@ -2639,8 +2663,12 @@ def open(self, path, **kwargs): return open_dataset(path, engine='pseudonetcdf', **kwargs) @contextlib.contextmanager - def roundtrip(self, data, save_kwargs={}, open_kwargs={}, + def roundtrip(self, data, save_kwargs=None, open_kwargs=None, allow_cleanup_failure=False): + if save_kwargs is None: + save_kwargs = {} + if open_kwargs is None: + open_kwargs = {} with create_tmp_file( allow_cleanup_failure=allow_cleanup_failure) as path: self.save(data, path, **save_kwargs) @@ -2847,10 +2875,14 @@ def create_tmp_geotiff(nx=4, ny=3, nz=3, transform_args=[5000, 80000, 1000, 2000.], crs={'units': 'm', 'no_defs': True, 'ellps': 'WGS84', 'proj': 'utm', 'zone': 18}, - open_kwargs={}): + open_kwargs=None): # yields a temporary geotiff file and a corresponding expected DataArray import rasterio from rasterio.transform import from_origin + + if open_kwargs is None: + open_kwargs = {} + with create_tmp_file(suffix='.tif', allow_cleanup_failure=ON_WINDOWS) as tmp_file: # allow 2d or 3d shapes