No public description

PiperOrigin-RevId: 573805160

xarray_beam/_src/core.py

@@ -329,7 +329,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/rechunk.py

@@ -295,7 +295,7 @@ def consolidate_fully(
         f'merging dataset chunks with variables {combined_vars} failed.\n'
         + repr_string
     ) from original_error
-  return core.Key(combined_offsets, combined_vars), dataset
+  return core.Key(combined_offsets, combined_vars), dataset  # pytype: disable=wrong-arg-types
 
 
 class _ConsolidateBase(beam.PTransform):
@@ -462,7 +462,7 @@ def split_variables(
   for var_name in dataset:
     new_dataset = dataset[[var_name]]
     offsets = {k: v for k, v in key.offsets.items() if k in new_dataset.dims}
-    new_key = core.Key(offsets, vars={var_name})
+    new_key = core.Key(offsets, vars={var_name})  # pytype: disable=wrong-arg-types
     yield new_key, new_dataset
 
 

xarray_beam/_src/zarr.py

@@ -143,11 +143,11 @@ def make_template(
 
 
 def _unchunked_vars(ds: xarray.Dataset) -> Set[str]:
-  return {k for k, v in ds.variables.items() if v.chunks is None}
+  return {k for k, v in ds.variables.items() if v.chunks is None}  # pytype: disable=bad-return-type
 
 
 def _chunked_vars(ds: xarray.Dataset) -> Set[str]:
-  return set(ds.variables.keys()) - _unchunked_vars(ds)
+  return set(ds.variables.keys()) - _unchunked_vars(ds)  # pytype: disable=bad-return-type
 
 
 def _make_template_from_chunked(dataset: xarray.Dataset) -> xarray.Dataset:
@@ -234,8 +234,8 @@ def _dask_to_zarr_chunksize(dim: str, sizes: Tuple[int, ...]) -> int:
 
 
 def _infer_zarr_chunks(dataset: xarray.Dataset) -> Dict[str, int]:
-  return {
-      dim: _dask_to_zarr_chunksize(dim, sizes)
+  return {  # pytype: disable=bad-return-type
+      dim: _dask_to_zarr_chunksize(dim, sizes)  # pytype: disable=wrong-arg-types
       for dim, sizes in dataset.chunks.items()
   }