diff --git a/tests/test_accessor.py b/tests/test_accessor.py new file mode 100644 index 0000000..7ec84d6 --- /dev/null +++ b/tests/test_accessor.py @@ -0,0 +1,17 @@ +import pytest +import xarray as xr + +import xarray_subset_grid.accessor # noqa: F401 + + +def test_data_vars_returns_empty_set_when_grid_not_recognized(): + ds = xr.Dataset( + data_vars={"foo": ("x", [1, 2, 3])}, + coords={"x": [0, 1, 2]}, + ) + + with pytest.warns(UserWarning, match="no grid type found in this dataset"): + accessor = ds.xsg + + assert accessor.grid is None + assert accessor.data_vars == set() diff --git a/xarray_subset_grid/accessor.py b/xarray_subset_grid/accessor.py index 7a4ca7a..1b127fd 100644 --- a/xarray_subset_grid/accessor.py +++ b/xarray_subset_grid/accessor.py @@ -1,4 +1,5 @@ import warnings +from typing import TYPE_CHECKING # from typing import Optional, Union import numpy as np @@ -15,6 +16,9 @@ UGrid, ) +if TYPE_CHECKING: + from xarray.core.coordinates import DatasetCoordinates + _grid_impls = [ FVCOMGrid, SELFEGrid, @@ -78,15 +82,13 @@ def data_vars(self) -> set[str]: data analysis. These can be discarded when subsetting the dataset when they are not needed. """ - if self._ds: + if self._grid: return self._grid.data_vars(self._ds) return set() @property - def coords(self) -> set[str]: - if self._ds: - return self._ds.coords - return set() + def coords(self) -> "DatasetCoordinates": + return self._ds.coords @property def grid_vars(self) -> set[str]: