Replace # pytype: disable suppressions with # pyrefly: ignore

xarray_beam/_src/core.py

@@ -457,11 +457,11 @@ def _first(self) -> xarray.Dataset:
   def _datasets(self) -> list[xarray.Dataset]:
     if isinstance(self.dataset, xarray.Dataset):
       return [self.dataset]
-    return list(self.dataset)  # pytype: disable=bad-return-type
+    return list(self.dataset)
 
   @functools.cached_property
   def expanded_chunks(self) -> dict[str, tuple[int, ...]]:
-    return normalize_expanded_chunks(self.chunks, self._first.sizes)  # pytype: disable=wrong-arg-types  # always-use-property-annotation
+    return normalize_expanded_chunks(self.chunks, self._first.sizes)  # pyrefly: ignore[bad-argument-type]
 
   @functools.cached_property
   def offsets(self) -> dict[str, list[int]]:
@@ -591,7 +591,7 @@ def sharded_dim(self) -> str | None:
         k: math.ceil(size / self.chunks.get(k, size))  # pyrefly: ignore[missing-attribute]
         for k, size in self._first.sizes.items()
     }
-    return max(lengths, key=lengths.get) if lengths else None  # pytype: disable=bad-return-type
+    return max(lengths, key=lengths.get) if lengths else None  # pyrefly: ignore[no-matching-overload]
 
   @functools.cached_property
   def shard_count(self) -> int | None:
@@ -615,7 +615,7 @@ def _iter_all_keys(self) -> Iterator[Key]:
         relevant_offsets = {
             k: v for k, v in self.offsets.items() if k in variable.dims
         }
-        yield from iter_chunk_keys(relevant_offsets, vars={name})  # pytype: disable=wrong-arg-types  # always-use-property-annotation
+        yield from iter_chunk_keys(relevant_offsets, vars={name})  # pyrefly: ignore[bad-argument-type]
 
   def _iter_shard_keys(
       self, shard_id: int | None, var_name: str | None
@@ -649,7 +649,7 @@ def _shard_inputs(self) -> list[tuple[int | None, str | None]]:
         inputs.extend([(i, name) for i in range(self.shard_count)])  # pyrefly: ignore[bad-argument-type]
       else:
         inputs.append((None, name))
-    return inputs  # pytype: disable=bad-return-type  # always-use-property-annotation
+    return inputs
 
   def expand(self, pcoll):
     if self.shard_count is None:
@@ -724,7 +724,7 @@ def _chunk_index_shapes(
         dims = tuple(d for d in variable.dims if d in self.expanded_chunks)
         shape = tuple(len(self.expanded_chunks[dim]) for dim in dims)
         out.append((name, dims, shape))
-    return out  # pytype: disable=bad-return-type
+    return out
 
   @functools.cached_property
   def _cumulative_sizes(self) -> np.ndarray:
@@ -743,7 +743,7 @@ def _index_to_key(self, position: int) -> Key:
     return Key(offsets, vars=None if name is None else {name})
 
   def _get_element(self, position: int) -> tuple[Key, xarray.Dataset]:
-    return self._key_to_chunks(self._index_to_key(position))  # pytype: disable=bad-return-type
+    return self._key_to_chunks(self._index_to_key(position))
 
   def expand(
       self, pbegin: beam.PBegin

xarray_beam/_src/dataset.py

@@ -162,7 +162,7 @@ def normalize_chunks(
     chunks = {k: chunks.get(k, default_chunks) for k in template.dims}  # pyrefly: ignore[bad-assignment, missing-attribute]
 
   defaults = previous_chunks if previous_chunks else template.sizes
-  chunks: dict[str, int | str] = {**defaults, **chunks}  # pytype: disable=annotation-type-mismatch
+  chunks: dict[str, int | str] = {**defaults, **chunks}  # pyrefly: ignore[invalid-argument]
 
   dtypes = {
       k: v.dtype for k, v in template.variables.items() if v.chunks is not None
@@ -502,7 +502,7 @@ def ptransform(self) -> beam.PTransform | beam.PCollection:
   @property
   def sizes(self) -> Mapping[str, int]:
     """Size of each dimension on this dataset."""
-    return dict(self.template.sizes)  # pytype: disable=bad-return-type
+    return dict(self.template.sizes)  # pyrefly: ignore[bad-return]
 
   @property
   def itemsize(self) -> int:
@@ -938,7 +938,7 @@ def map_blocks(
           old_sizes=self.sizes,
           old_chunks=self.chunks,
           new_sizes=template.sizes,  # pyrefly: ignore[bad-argument-type]
-      )  # pytype: disable=wrong-arg-types
+      )
 
     for dim, old_chunks in self.chunks.items():
       if old_chunks < self.sizes[dim]:

xarray_beam/_src/rechunk.py

@@ -286,7 +286,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  # pytype: disable=wrong-arg-types
+  return core.Key(combined_offsets, combined_vars), dataset
 
 
 def _round_chunk_key(
@@ -465,7 +465,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})  # pytype: disable=wrong-arg-types
+    new_key = core.Key(offsets, vars={var_name})  # pyrefly: ignore[bad-argument-type]
     yield new_key, new_dataset
 
 

xarray_beam/_src/zarr.py

@@ -227,11 +227,11 @@ def replace_template_dims(
 
 
 def _unchunked_vars(ds: xarray.Dataset) -> set[str]:
-  return {k for k, v in ds.variables.items() if v.chunks is None}  # pytype: disable=bad-return-type
+  return {k for k, v in ds.variables.items() if v.chunks is None}  # pyrefly: ignore[bad-return]
 
 
 def _chunked_vars(ds: xarray.Dataset) -> set[str]:
-  return set(ds.variables.keys()) - _unchunked_vars(ds)  # pytype: disable=bad-return-type
+  return set(ds.variables.keys()) - _unchunked_vars(ds)  # pyrefly: ignore[bad-return]
 
 
 def _make_template_from_chunked(dataset: xarray.Dataset) -> xarray.Dataset:
@@ -298,8 +298,8 @@ def _dask_to_zarr_chunksize(dim: str, sizes: tuple[int, ...]) -> int:
 
 
 def _zarr_from_dask_chunks(dataset: xarray.Dataset) -> dict[str, int]:
-  return {  # pytype: disable=bad-return-type
-      dim: _dask_to_zarr_chunksize(dim, sizes)  # pytype: disable=wrong-arg-types
+  return {  # pyrefly: ignore[bad-return]
+      dim: _dask_to_zarr_chunksize(dim, sizes)  # pyrefly: ignore[bad-argument-type]
       for dim, sizes in dataset.chunks.items()
   }