diff --git a/virtualizarr/backend.py b/virtualizarr/backend.py index 0322f60..32403d0 100644 --- a/virtualizarr/backend.py +++ b/virtualizarr/backend.py @@ -7,8 +7,7 @@ Optional, ) -from xarray import Dataset -from xarray.core.indexes import Index +from xarray import Dataset, Index from virtualizarr.manifests import ManifestArray from virtualizarr.readers import ( diff --git a/virtualizarr/readers/common.py b/virtualizarr/readers/common.py index 54aedfe..9be2b45 100644 --- a/virtualizarr/readers/common.py +++ b/virtualizarr/readers/common.py @@ -11,11 +11,16 @@ cast, ) -import xarray as xr -from xarray import Dataset +from xarray import ( + Coordinates, + Dataset, + Index, + IndexVariable, + Variable, + open_dataset, +) from xarray.backends import AbstractDataStore, BackendArray -from xarray.core.indexes import Index, PandasIndex -from xarray.core.variable import IndexVariable, Variable +from xarray.core.indexes import PandasIndex from virtualizarr.manifests import ManifestArray from virtualizarr.utils import _FsspecFSFromFilepath @@ -62,7 +67,7 @@ def open_loadable_vars_and_indexes( # fpath can be `Any` thanks to fsspec.filesystem(...).open() returning Any. # We'll (hopefully safely) cast it to what xarray is expecting, but this might let errors through. - ds = xr.open_dataset( + ds = open_dataset( cast(XArrayOpenT, fpath), drop_variables=drop_variables, group=group, @@ -113,7 +118,7 @@ def construct_virtual_dataset( data_vars, coords = separate_coords(vars, indexes, coord_names) - vds = xr.Dataset( + vds = Dataset( data_vars, coords=coords, # indexes={}, # TODO should be added in a later version of xarray @@ -126,10 +131,10 @@ def construct_virtual_dataset( def separate_coords( - vars: Mapping[str, xr.Variable], + vars: Mapping[str, Variable], indexes: MutableMapping[str, Index], coord_names: Iterable[str] | None = None, -) -> tuple[dict[str, xr.Variable], xr.Coordinates]: +) -> tuple[dict[str, Variable], Coordinates]: """ Try to generate a set of coordinates that won't cause xarray to automatically build a pandas.Index for the 1D coordinates. @@ -144,7 +149,7 @@ def separate_coords( # split data and coordinate variables (promote dimension coordinates) data_vars = {} coord_vars: dict[ - str, tuple[Hashable, Any, dict[Any, Any], dict[Any, Any]] | xr.Variable + str, tuple[Hashable, Any, dict[Any, Any], dict[Any, Any]] | Variable ] = {} for name, var in vars.items(): if name in coord_names or var.dims == (name,): @@ -164,7 +169,7 @@ def separate_coords( else: data_vars[name] = var - coords = xr.Coordinates(coord_vars, indexes=indexes) + coords = Coordinates(coord_vars, indexes=indexes) return data_vars, coords diff --git a/virtualizarr/readers/dmrpp.py b/virtualizarr/readers/dmrpp.py index 766b1c6..c9095c7 100644 --- a/virtualizarr/readers/dmrpp.py +++ b/virtualizarr/readers/dmrpp.py @@ -6,9 +6,7 @@ from xml.etree import ElementTree as ET import numpy as np -from xarray import Coordinates, Dataset -from xarray.core.indexes import Index -from xarray.core.variable import Variable +from xarray import Coordinates, Dataset, Index, Variable from virtualizarr.manifests import ChunkManifest, ManifestArray from virtualizarr.readers.common import VirtualBackend diff --git a/virtualizarr/readers/fits.py b/virtualizarr/readers/fits.py index 618d81c..de93bc1 100644 --- a/virtualizarr/readers/fits.py +++ b/virtualizarr/readers/fits.py @@ -1,7 +1,6 @@ from typing import Iterable, Mapping, Optional -from xarray import Dataset -from xarray.core.indexes import Index +from xarray import Dataset, Index from virtualizarr.readers.common import ( VirtualBackend, diff --git a/virtualizarr/readers/hdf5.py b/virtualizarr/readers/hdf5.py index c0d38e2..91e5b6f 100644 --- a/virtualizarr/readers/hdf5.py +++ b/virtualizarr/readers/hdf5.py @@ -1,7 +1,6 @@ from typing import Iterable, Mapping, Optional -from xarray import Dataset -from xarray.core.indexes import Index +from xarray import Dataset, Index from virtualizarr.readers.common import ( VirtualBackend, diff --git a/virtualizarr/readers/kerchunk.py b/virtualizarr/readers/kerchunk.py index 35fa493..2f1ff4b 100644 --- a/virtualizarr/readers/kerchunk.py +++ b/virtualizarr/readers/kerchunk.py @@ -1,8 +1,7 @@ from typing import Iterable, Mapping, Optional import ujson -from xarray import Dataset -from xarray.core.indexes import Index +from xarray import Dataset, Index from virtualizarr.readers.common import VirtualBackend from virtualizarr.translators.kerchunk import dataset_from_kerchunk_refs diff --git a/virtualizarr/readers/netcdf3.py b/virtualizarr/readers/netcdf3.py index 30c6746..25f212c 100644 --- a/virtualizarr/readers/netcdf3.py +++ b/virtualizarr/readers/netcdf3.py @@ -1,7 +1,6 @@ from typing import Iterable, Mapping, Optional -from xarray import Dataset -from xarray.core.indexes import Index +from xarray import Dataset, Index from virtualizarr.readers.common import ( VirtualBackend, diff --git a/virtualizarr/readers/tiff.py b/virtualizarr/readers/tiff.py index bb32e64..d9c440b 100644 --- a/virtualizarr/readers/tiff.py +++ b/virtualizarr/readers/tiff.py @@ -1,8 +1,7 @@ import warnings from typing import Iterable, Mapping, Optional -from xarray import Dataset -from xarray.core.indexes import Index +from xarray import Dataset, Index from virtualizarr.readers.common import ( VirtualBackend, diff --git a/virtualizarr/readers/zarr_v3.py b/virtualizarr/readers/zarr_v3.py index a1f4ab7..4a867ff 100644 --- a/virtualizarr/readers/zarr_v3.py +++ b/virtualizarr/readers/zarr_v3.py @@ -4,9 +4,7 @@ import numcodecs import numpy as np -from xarray import Dataset -from xarray.core.indexes import Index -from xarray.core.variable import Variable +from xarray import Dataset, Index, Variable from virtualizarr.manifests import ChunkManifest, ManifestArray from virtualizarr.readers.common import VirtualBackend, separate_coords