No public description

PiperOrigin-RevId: 573805160

xarray_beam/_src/core.py

@@ -18,6 +18,7 @@
     AbstractSet,
     Dict,
     Generic,
+    Hashable,
     Iterator,
     List,
     Mapping,
@@ -92,7 +93,7 @@ class Key:
   def __init__(
       self,
       offsets: Optional[Mapping[str, int]] = None,
-      vars: Optional[AbstractSet[str]] = None,
+      vars: Optional[AbstractSet[Hashable]] = None,
   ):
     if offsets is None:
       offsets = {}
@@ -329,7 +330,7 @@ def _first(self) -> xarray.Dataset:
   def _datasets(self) -> List[xarray.Dataset]:
     if isinstance(self.dataset, xarray.Dataset):
       return [self.dataset]
-    return list(self.dataset)
+    return list(self.dataset)  # pytype: disable=bad-return-type
 
   def _validate(self, dataset, split_vars):
     """Raise errors if input parameters are invalid."""

xarray_beam/_src/dataset.py

@@ -107,7 +107,7 @@ def collect_with_direct_runner(self) -> xarray.Dataset:
   @property
   def sizes(self) -> dict[str, int]:
     """Size of each dimension on this dataset."""
-    return dict(self.template.sizes)
+    return dict(self.template.sizes)  # pytype: disable=bad-return-type
 
   def pipe(self, func, *args, **kwargs):
     return func(*args, **kwargs)

xarray_beam/_src/zarr.py

@@ -20,6 +20,7 @@
     Any,
     AbstractSet,
     Dict,
+    Hashable,
     Optional,
     Mapping,
     Set,
@@ -43,7 +44,7 @@
 WritableStore = Union[str, MutableMapping[str, bytes]]
 
 
-def _infer_chunks(dataset: xarray.Dataset) -> Mapping[str, int]:
+def _infer_chunks(dataset: xarray.Dataset) -> Mapping[Hashable, int]:
   """Infer chunks for an xarray.Dataset loaded from Zarr."""
   # The original Zarr array chunks (as tuples) are stored in the "encoding"
   # dictionary on each DataArray object.
@@ -68,7 +69,7 @@ def _infer_chunks(dataset: xarray.Dataset) -> Mapping[str, int]:
 
 def open_zarr(
     store: ReadableStore, **kwargs: Any
-) -> Tuple[xarray.Dataset, Mapping[str, int]]:
+) -> Tuple[xarray.Dataset, Mapping[Hashable, int]]:
   """Returns a lazily indexable xarray.Dataset and chunks from a Zarr store.
 
   Only Zarr stores with the consistent chunking between non-indexed variables
@@ -102,7 +103,7 @@ def _raise_template_error():
 
 def make_template(
     dataset: xarray.Dataset,
-    lazy_vars: Optional[AbstractSet[str]] = None,
+    lazy_vars: Optional[AbstractSet[Hashable]] = None,
 ) -> xarray.Dataset:
   """Make a lazy Dask xarray.Dataset for use only as a template.
 
@@ -142,11 +143,11 @@ def make_template(
   return result
 
 
-def _unchunked_vars(ds: xarray.Dataset) -> Set[str]:
+def _unchunked_vars(ds: xarray.Dataset) -> Set[Hashable]:
   return {k for k, v in ds.variables.items() if v.chunks is None}
 
 
-def _chunked_vars(ds: xarray.Dataset) -> Set[str]:
+def _chunked_vars(ds: xarray.Dataset) -> Set[Hashable]:
   return set(ds.variables.keys()) - _unchunked_vars(ds)
 
 
@@ -201,7 +202,7 @@ def _verify_template_is_lazy(template: xarray.Dataset):
 
 def _override_chunks(
     dataset: xarray.Dataset,
-    chunks: Mapping[str, int],
+    chunks: Mapping[Hashable, int],
 ) -> xarray.Dataset:
   """Override chunks on a Dataset, for already chunked variables only."""
 
@@ -219,7 +220,7 @@ def maybe_rechunk(variable):
   return xarray.Dataset(data_vars, coords, dataset.attrs)
 
 
-def _dask_to_zarr_chunksize(dim: str, sizes: Tuple[int, ...]) -> int:
+def _dask_to_zarr_chunksize(dim: Hashable, sizes: Tuple[int, ...]) -> int:
   if not sizes:
     return 0
   # It's OK for the last chunk of Zarr array to have smaller size. Otherwise,
@@ -233,7 +234,7 @@ def _dask_to_zarr_chunksize(dim: str, sizes: Tuple[int, ...]) -> int:
   return sizes[0]
 
 
-def _infer_zarr_chunks(dataset: xarray.Dataset) -> Dict[str, int]:
+def _infer_zarr_chunks(dataset: xarray.Dataset) -> Dict[Hashable, int]:
   return {
       dim: _dask_to_zarr_chunksize(dim, sizes)
       for dim, sizes in dataset.chunks.items()