diff --git a/xarray_beam/_src/core.py b/xarray_beam/_src/core.py index 4eec0e6..a4be8a5 100644 --- a/xarray_beam/_src/core.py +++ b/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 diff --git a/xarray_beam/_src/dataset.py b/xarray_beam/_src/dataset.py index d15f198..3b0a051 100644 --- a/xarray_beam/_src/dataset.py +++ b/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]: diff --git a/xarray_beam/_src/rechunk.py b/xarray_beam/_src/rechunk.py index 56be385..bd85ecf 100644 --- a/xarray_beam/_src/rechunk.py +++ b/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 diff --git a/xarray_beam/_src/zarr.py b/xarray_beam/_src/zarr.py index 1758192..56d0f63 100644 --- a/xarray_beam/_src/zarr.py +++ b/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() }