From 8a811f22c8d562317f1c009c0a7b7cbb2023fc52 Mon Sep 17 00:00:00 2001 From: Hana Joo Date: Wed, 8 Jul 2026 01:48:28 -0700 Subject: [PATCH] Adding type suppressions for pyrefly PiperOrigin-RevId: 944352684 --- docs/conf.py | 2 +- examples/xbeam_rechunk.py | 4 +- xarray_beam/_src/core.py | 28 +++++++------- xarray_beam/_src/dataset.py | 73 +++++++++++++++++++------------------ xarray_beam/_src/rechunk.py | 8 ++-- xarray_beam/_src/zarr.py | 22 +++++------ 6 files changed, 69 insertions(+), 68 deletions(-) diff --git a/docs/conf.py b/docs/conf.py index fe705e6..0b5a966 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -118,7 +118,7 @@ def linkcode_resolve(domain, info): except Exception as e: print(f'did not find source code for: {info}: {e}') return None - filename = os.path.relpath( + filename = os.path.relpath( # pyrefly: ignore[no-matching-overload] filename, start=os.path.dirname(xarray_beam.__file__) ) lines = f'#L{linenum}-L{linenum + len(source)}' if linenum else '' diff --git a/examples/xbeam_rechunk.py b/examples/xbeam_rechunk.py index cf0d023..9bebb75 100644 --- a/examples/xbeam_rechunk.py +++ b/examples/xbeam_rechunk.py @@ -88,10 +88,10 @@ def main(argv): with beam.Pipeline(runner=RUNNER.value, argv=argv) as root: root |= ( - xbeam.Dataset.from_zarr(INPUT_PATH.value, split_vars=True) + xbeam.Dataset.from_zarr(INPUT_PATH.value, split_vars=True) # pyrefly: ignore[bad-argument-type] .rechunk(target_chunks if target_shards is None else target_shards) .to_zarr( - OUTPUT_PATH.value, + OUTPUT_PATH.value, # pyrefly: ignore[bad-argument-type] zarr_chunks=target_chunks, zarr_shards=target_shards, zarr_format=ZARR_FORMAT.value, diff --git a/xarray_beam/_src/core.py b/xarray_beam/_src/core.py index 3aa80d9..4eec0e6 100644 --- a/xarray_beam/_src/core.py +++ b/xarray_beam/_src/core.py @@ -158,7 +158,7 @@ def replace( vars = self.vars if indices is _DEFAULT: indices = self.indices - return type(self)(offsets, vars, indices) + return type(self)(offsets, vars, indices) # pyrefly: ignore[bad-argument-type] def with_offsets(self, **offsets: int | None) -> Key: """Replace some offsets with new values. @@ -401,7 +401,7 @@ def normalize_expanded_chunks( if dim not in chunks or chunks[dim] == -1: result[dim] = (dim_size,) elif isinstance(chunks[dim], tuple): - total = sum(chunks[dim]) + total = sum(chunks[dim]) # pyrefly: ignore[no-matching-overload] if total != dim_size: raise ValueError( f"sum of provided chunks does not match size of dimension {dim}: " @@ -409,7 +409,7 @@ def normalize_expanded_chunks( ) result[dim] = chunks[dim] else: - multiple, remainder = divmod(dim_size, chunks[dim]) + multiple, remainder = divmod(dim_size, chunks[dim]) # pyrefly: ignore[no-matching-overload] result[dim] = multiple * (chunks[dim],) + ( (remainder,) if remainder else () ) @@ -439,9 +439,9 @@ def __init__( dask_chunks = self._first.chunks if not dask_chunks: raise ValueError("dataset must be chunked or chunks must be provided") - chunks = dask_to_xbeam_chunks(dask_chunks) + chunks = dask_to_xbeam_chunks(dask_chunks) # pyrefly: ignore[bad-assignment] - for k in chunks: + for k in chunks: # pyrefly: ignore[not-iterable] if k not in self._first.dims: raise ValueError( f"chunks key {k!r} is not a dimension on the provided dataset(s)" @@ -538,7 +538,7 @@ def _key_to_chunks(self, key: Key) -> tuple[Key, DatasetOrDatasets]: if isinstance(self.dataset, xarray.Dataset): return key, results[0] else: - return key, results + return key, results # pyrefly: ignore[bad-return] @export @@ -588,7 +588,7 @@ def sharded_dim(self) -> str | None: # We use the simple heuristic of only sharding inputs along the dimension # with the most chunks. lengths = { - k: math.ceil(size / self.chunks.get(k, size)) + 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 @@ -624,29 +624,29 @@ def _iter_shard_keys( if var_name is None: offsets = self.offsets else: - offsets = {dim: self.offsets[dim] for dim in self._first[var_name].dims} + offsets = {dim: self.offsets[dim] for dim in self._first[var_name].dims} # pyrefly: ignore[bad-index] if shard_id is None: assert self.split_vars - yield from iter_chunk_keys(offsets, vars={var_name}) + yield from iter_chunk_keys(offsets, vars={var_name}) # pyrefly: ignore[bad-argument-type] else: assert self.split_vars == (var_name is not None) dim = self.sharded_dim - count = math.ceil(len(self.offsets[dim]) / self.shard_count) + count = math.ceil(len(self.offsets[dim]) / self.shard_count) # pyrefly: ignore[bad-index, unsupported-operation] dim_slice = slice(shard_id * count, (shard_id + 1) * count) - offsets = {**offsets, dim: offsets[dim][dim_slice]} + offsets = {**offsets, dim: offsets[dim][dim_slice]} # pyrefly: ignore[bad-index, invalid-argument] vars_ = {var_name} if self.split_vars else None - yield from iter_chunk_keys(offsets, vars=vars_) + yield from iter_chunk_keys(offsets, vars=vars_) # pyrefly: ignore[bad-argument-type] def _shard_inputs(self) -> list[tuple[int | None, str | None]]: """Create inputs for sharded key iterators.""" if not self.split_vars: - return [(i, None) for i in range(self.shard_count)] + return [(i, None) for i in range(self.shard_count)] # pyrefly: ignore[bad-argument-type] inputs = [] for name, variable in self._first.items(): if self.sharded_dim in variable.dims: - inputs.extend([(i, name) for i in range(self.shard_count)]) + 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 diff --git a/xarray_beam/_src/dataset.py b/xarray_beam/_src/dataset.py index 81ead5d..396819e 100644 --- a/xarray_beam/_src/dataset.py +++ b/xarray_beam/_src/dataset.py @@ -72,7 +72,7 @@ def _to_human_size(nbytes: int) -> str: for unit in ['B', 'kB', 'MB', 'GB', 'TB', 'PB', 'EB']: if nbytes < 1000: return f'{_at_least_two_digits(nbytes)}{unit}' - nbytes /= 1000 + nbytes /= 1000 # pyrefly: ignore[bad-assignment] nbytes *= 1000 return f'{_at_least_two_digits(nbytes)}EB' @@ -141,7 +141,7 @@ def normalize_chunks( "chunks='auto'. Supply an explicit number of bytes instead, e.g., " "chunks='100MB'." ) - chunks = {k: chunks for k in template.dims} + chunks = {k: chunks for k in template.dims} # pyrefly: ignore[bad-assignment] elif isinstance(chunks, Mapping): string_chunks = {v for v in chunks.values() if isinstance(v, str)} if len(string_chunks) > 1: @@ -157,9 +157,9 @@ def normalize_chunks( else: raise TypeError(f'chunks must be a string or a mapping, got {chunks=}') - if ... in chunks: - default_chunks = chunks[...] - chunks = {k: chunks.get(k, default_chunks) for k in template.dims} + if ... in chunks: # pyrefly: ignore[not-iterable, unsupported-operation] + default_chunks = chunks[...] # pyrefly: ignore[bad-index] + chunks = {k: chunks.get(k, default_chunks) for k in template.dims} # pyrefly: ignore[bad-assignment] defaults = previous_chunks if previous_chunks else template.sizes chunks: dict[str, int | str] = {**defaults, **chunks} # pytype: disable=annotation-type-mismatch @@ -253,6 +253,7 @@ def _normalize_and_validate_chunk( key = key.replace(vars=set(dataset.keys())) elif key.vars != set(dataset.keys()): raise ValueError( + # pyrefly: ignore[bad-specialization] f'dataset keys {sorted(dataset.keys())} do not match' f' key.vars={sorted(key.vars)}' ) @@ -262,7 +263,7 @@ def _normalize_and_validate_chunk( new_offsets = dict(key.offsets) for dim in dataset.dims: if dim not in new_offsets: - new_offsets[dim] = 0 + new_offsets[dim] = 0 # pyrefly: ignore[unsupported-operation] if len(new_offsets) != len(key.offsets): key = key.replace(offsets=new_offsets) @@ -353,7 +354,7 @@ def _apply_to_each_chunk( key.offsets.get(dim, 0) // old_chunks.get(dim, 1) * new_chunks[dim] ) new_vars = set(new_chunk) if key.vars is not None else None - new_key = core.Key(new_offsets, new_vars) + new_key = core.Key(new_offsets, new_vars) # pyrefly: ignore[bad-argument-type] return new_key, new_chunk @@ -516,7 +517,7 @@ def itemsize(self) -> int: def bytes_per_chunk(self) -> int: """Estimate of the number of bytes per dataset chunk.""" variable_sizes = [ - v.dtype.itemsize * math.prod(self.chunks[d] for d in v.dims) + v.dtype.itemsize * math.prod(self.chunks[d] for d in v.dims) # pyrefly: ignore[bad-index] for v in self.template.values() ] return max(variable_sizes) if self.split_vars else sum(variable_sizes) @@ -528,7 +529,7 @@ def chunk_count(self) -> int: total = 0 for variable in self.template.values(): total += math.prod( - math.ceil(self.sizes[d] / self.chunks[d]) for d in variable.dims + math.ceil(self.sizes[d] / self.chunks[d]) for d in variable.dims # pyrefly: ignore[bad-index] ) return total else: @@ -612,13 +613,13 @@ def from_ptransform( f' got {chunks}' ) - chunks = normalize_chunks(chunks, template) + chunks = normalize_chunks(chunks, template) # pyrefly: ignore[bad-assignment] ptransform = ptransform | label >> beam.MapTuple( functools.partial( _normalize_and_validate_chunk, template, chunks, split_vars ) ) - return cls(template, chunks, split_vars, ptransform) + return cls(template, chunks, split_vars, ptransform) # pyrefly: ignore[bad-argument-type] @classmethod def from_xarray( @@ -650,13 +651,13 @@ def from_xarray( label = _get_label('from_xarray') template = zarr.make_template(source) if previous_chunks is None: - previous_chunks = source.sizes - chunks = normalize_chunks(chunks, template, split_vars, previous_chunks) - ptransform = core.DatasetToChunks(source, chunks, split_vars) + previous_chunks = source.sizes # pyrefly: ignore[bad-assignment] + chunks = normalize_chunks(chunks, template, split_vars, previous_chunks) # pyrefly: ignore[bad-assignment] + ptransform = core.DatasetToChunks(source, chunks, split_vars) # pyrefly: ignore[bad-argument-type] ptransform.label = label if pipeline is not None: ptransform = _LazyPCollection(pipeline, ptransform) - return cls(template, dict(chunks), split_vars, ptransform) + return cls(template, dict(chunks), split_vars, ptransform) # pyrefly: ignore[no-matching-overload] @classmethod def from_zarr( @@ -689,10 +690,10 @@ def from_zarr( label = _get_label('from_zarr') source, previous_chunks = zarr.open_zarr(path) if chunks is None: - chunks = previous_chunks + chunks = previous_chunks # pyrefly: ignore[bad-assignment] result = cls.from_xarray( source, - chunks, + chunks, # pyrefly: ignore[bad-argument-type] split_vars=split_vars, previous_chunks=previous_chunks, ) @@ -709,7 +710,7 @@ def _zarr_chunks_per_shard_to_chunks( """Convert chunks per shard to chunks.""" chunks_per_shard = dict(zarr_chunks_per_shard) if ... in chunks_per_shard: - default_cps = chunks_per_shard.pop(...) + default_cps = chunks_per_shard.pop(...) # pyrefly: ignore[bad-argument-type] else: default_cps = 1 @@ -808,7 +809,7 @@ def to_zarr( label = _get_label('to_zarr') if zarr_shards is not None: - zarr_shards = normalize_chunks( + zarr_shards = normalize_chunks( # pyrefly: ignore[bad-assignment] zarr_shards, self.template, split_vars=self.split_vars, @@ -821,9 +822,9 @@ def to_zarr( 'cannot supply both zarr_chunks_per_shard and zarr_chunks' ) if zarr_shards is None: - zarr_shards = self.chunks - zarr_chunks = self._zarr_chunks_per_shard_to_chunks( - zarr_chunks_per_shard, zarr_shards + zarr_shards = self.chunks # pyrefly: ignore[bad-assignment] + zarr_chunks = self._zarr_chunks_per_shard_to_chunks( # pyrefly: ignore[bad-assignment] + zarr_chunks_per_shard, zarr_shards # pyrefly: ignore[bad-argument-type] ) if zarr_chunks is None: @@ -831,7 +832,7 @@ def to_zarr( raise ValueError('cannot supply zarr_shards without zarr_chunks') zarr_chunks = {} - zarr_chunks = normalize_chunks( + zarr_chunks = normalize_chunks( # pyrefly: ignore[bad-assignment] zarr_chunks, self.template, split_vars=self.split_vars, @@ -842,15 +843,15 @@ def to_zarr( # chunk sizes, which means shard sizes must be rounded up to be larger # than the full dimension size. This will likely be relaxed in the future: # https://github.com/zarr-developers/zarr-extensions/issues/34 - zarr_shards = dict(zarr_shards) + zarr_shards = dict(zarr_shards) # pyrefly: ignore[no-matching-overload] for k in zarr_shards: if zarr_shards[k] == self.sizes[k]: zarr_shards[k] = ( - math.ceil(zarr_shards[k] / zarr_chunks[k]) * zarr_chunks[k] + math.ceil(zarr_shards[k] / zarr_chunks[k]) * zarr_chunks[k] # pyrefly: ignore[bad-index, unsupported-operation] ) self._check_shards_or_chunks(zarr_shards, 'shards') else: - self._check_shards_or_chunks(zarr_chunks, 'chunks') + self._check_shards_or_chunks(zarr_chunks, 'chunks') # pyrefly: ignore[bad-argument-type] if zarr_shards is not None and zarr_format is None: zarr_format = 3 # required for shards @@ -858,7 +859,7 @@ def to_zarr( return self.ptransform | label >> zarr.ChunksToZarr( path, self.template, - zarr_chunks=zarr_chunks, + zarr_chunks=zarr_chunks, # pyrefly: ignore[bad-argument-type] zarr_shards=zarr_shards, zarr_format=zarr_format, stage_locally=stage_locally, @@ -936,7 +937,7 @@ def map_blocks( chunks = _infer_new_chunks( old_sizes=self.sizes, old_chunks=self.chunks, - new_sizes=template.sizes, + new_sizes=template.sizes, # pyrefly: ignore[bad-argument-type] ) # pytype: disable=wrong-arg-types for dim, old_chunks in self.chunks.items(): @@ -998,7 +999,7 @@ def rechunk( if split_vars is None: split_vars = self.split_vars - chunks = normalize_chunks( + chunks = normalize_chunks( # pyrefly: ignore[bad-assignment] chunks, self.template, split_vars=split_vars, @@ -1007,15 +1008,15 @@ def rechunk( pipeline, ptransform = _split_lazy_pcollection(self._ptransform) if isinstance(ptransform, core.DatasetToChunks) and all( - chunks[k] % self.chunks[k] == 0 for k in chunks + chunks[k] % self.chunks[k] == 0 for k in chunks # pyrefly: ignore[bad-index, not-iterable] ): # Rechunking can be performed by re-reading the source dataset with new # chunks, rather than using a separate rechunking transform. - ptransform = core.DatasetToChunks(ptransform.dataset, chunks, split_vars) + ptransform = core.DatasetToChunks(ptransform.dataset, chunks, split_vars) # pyrefly: ignore[bad-argument-type] ptransform.label = _concat_labels(ptransform.label, label) if pipeline is not None: ptransform = _LazyPCollection(pipeline, ptransform) - return type(self)(self.template, chunks, split_vars, ptransform) + return type(self)(self.template, chunks, split_vars, ptransform) # pyrefly: ignore[bad-argument-type] # Need to do a full rechunking. # If also splitting variables, do that first because smaller itemsize allows @@ -1024,14 +1025,14 @@ def rechunk( rechunk_transform = rechunk.Rechunk( prechunked.sizes, prechunked.chunks, - chunks, + chunks, # pyrefly: ignore[bad-argument-type] itemsize=prechunked.itemsize, min_mem=min_mem, max_mem=max_mem, ) ptransform = prechunked.ptransform | label >> rechunk_transform rechunked = type(self)( - self.template, chunks, prechunked.split_vars, ptransform + self.template, chunks, prechunked.split_vars, ptransform # pyrefly: ignore[bad-argument-type] ) result = rechunked if split_vars else rechunked.consolidate_variables() return result @@ -1086,13 +1087,13 @@ def mean( else: dims = dim if label is None: - label = _get_label(f"mean_{'_'.join(dims)}") + label = _get_label(f"mean_{'_'.join(dims)}") # pyrefly: ignore[no-matching-overload] template = zarr.make_template( self.template.mean(dim=dims, skipna=skipna, dtype=dtype) ) new_chunks = {k: v for k, v in self.chunks.items() if k not in dims} ptransform = self.ptransform | label >> combiners.MultiStageMean( - dims=dims, + dims=dims, # pyrefly: ignore[bad-argument-type] skipna=skipna, dtype=dtype, chunks=self.chunks, diff --git a/xarray_beam/_src/rechunk.py b/xarray_beam/_src/rechunk.py index d30d767..56be385 100644 --- a/xarray_beam/_src/rechunk.py +++ b/xarray_beam/_src/rechunk.py @@ -47,7 +47,7 @@ def normalize_chunks( if dim not in chunks or chunks[dim] == -1: new_chunk_size = size elif isinstance(chunks[dim], tuple): - unique_chunks = set(chunks[dim]) + unique_chunks = set(chunks[dim]) # pyrefly: ignore[bad-argument-type] if len(unique_chunks) != 1: raise ValueError( f'chunks for dimension {dim} are not constant: {unique_chunks}', @@ -139,7 +139,7 @@ def _consolidate_chunks_in_var_group( kwargs.update(combine_kwargs) try: - combined_dataset = xarray.combine_nested( + combined_dataset = xarray.combine_nested( # pyrefly: ignore[no-matching-overload] nested_array.tolist(), concat_dim=list(offsets), **kwargs ) return combined_key, combined_dataset @@ -225,7 +225,7 @@ def consolidate_variables( key = core.Key(offsets, new_vars) try: - dataset = xarray.merge(chunks, **kwargs) + dataset = xarray.merge(chunks, **kwargs) # pyrefly: ignore[no-matching-overload] except (ValueError, xarray.MergeError) as original_error: repr_string = '\n'.join(repr(ds) for ds in chunks[:2]) if len(chunks) > 2: @@ -276,7 +276,7 @@ def consolidate_fully( kwargs.update(merge_kwargs) try: - dataset = xarray.merge(concatenated_chunks, **kwargs) + dataset = xarray.merge(concatenated_chunks, **kwargs) # pyrefly: ignore[no-matching-overload] except (ValueError, xarray.MergeError) as original_error: repr_string = '\n'.join(repr(ds) for ds in concatenated_chunks[:2]) if len(concatenated_chunks) > 2: diff --git a/xarray_beam/_src/zarr.py b/xarray_beam/_src/zarr.py index 60775c2..61713f4 100644 --- a/xarray_beam/_src/zarr.py +++ b/xarray_beam/_src/zarr.py @@ -124,18 +124,18 @@ def make_template( Non-lazy variables are loaded in memory as NumPy arrays. """ if lazy_vars is None: - lazy_vars = set(dataset.keys()) + lazy_vars = set(dataset.keys()) # pyrefly: ignore[bad-assignment] lazy_vars.update(k for k in dataset.coords if k not in dataset.indexes) result = dataset.copy() # load non-lazy variables into memory - result.update(dataset.drop_vars(lazy_vars).compute()) + result.update(dataset.drop_vars(lazy_vars).compute()) # pyrefly: ignore[bad-argument-type] # override the lazy variables delayed = dask.delayed(_raise_template_error)() for k, v in dataset.variables.items(): - if k in lazy_vars: + if k in lazy_vars: # pyrefly: ignore[not-iterable] # names of dask arrays are used for keeping track of results, so arrays # with the same name cannot have different shape or dtype name = f"make_template_{'x'.join(map(str, v.shape))}_{v.dtype}" @@ -342,7 +342,7 @@ def _finalize_setup_zarr_args( template = _make_template_from_chunked(template) chunks = _finalize_chunks(template, chunks) if shards is not None: - shards = _finalize_chunks(template, chunks | shards) + shards = _finalize_chunks(template, chunks | shards) # pyrefly: ignore[unsupported-operation] if not all(shards[k] % chunks[k] == 0 for k in chunks): # raise a better error message than the user would see from zarr-python raise ValueError( @@ -362,11 +362,11 @@ def _get_chunk_and_shard_encoding( for var_name in _chunked_vars(template): assert isinstance(var_name, str) variable = template.variables[var_name] - chunks = tuple(zarr_chunks[dim] for dim in variable.dims) + chunks = tuple(zarr_chunks[dim] for dim in variable.dims) # pyrefly: ignore[bad-index] encoding[var_name] = {'chunks': chunks} if zarr_shards is not None: encoding[var_name]['shards'] = tuple( - zarr_shards[dim] for dim in variable.dims + zarr_shards[dim] for dim in variable.dims # pyrefly: ignore[bad-index] ) return encoding @@ -411,10 +411,10 @@ def _setup_zarr( del var.encoding['chunks'] chunk_encoding = _get_chunk_and_shard_encoding( - template, zarr_chunks, zarr_shards + template, zarr_chunks, zarr_shards # pyrefly: ignore[bad-argument-type] ) - encoding = { - k: encoding.get(k, {}) | chunk_encoding.get(k, {}) + encoding = { # pyrefly: ignore[bad-assignment] + k: encoding.get(k, {}) | chunk_encoding.get(k, {}) # pyrefly: ignore[no-matching-overload] for k in template.variables } encoding_str = pprint.pformat(encoding, sort_dicts=False) @@ -543,7 +543,7 @@ def validate_zarr_chunk( f'{unexpected_indexes}' ) - region = core.offsets_to_slices(key.offsets, chunk.sizes) + region = core.offsets_to_slices(key.offsets, chunk.sizes) # pyrefly: ignore[bad-argument-type] for dim, full_index in template.indexes.items(): if dim in chunk.indexes: expected_index = full_index[region[dim]] @@ -610,7 +610,7 @@ def write_chunk_to_zarr( k for k in chunk.variables if k in _unchunked_vars(template) ] writable_chunk = chunk.drop_vars(already_written) - region = core.offsets_to_slices(key.offsets, writable_chunk.sizes) + region = core.offsets_to_slices(key.offsets, writable_chunk.sizes) # pyrefly: ignore[bad-argument-type] # Ensure the arrays in writable_chunk are each stored in a single dask chunk. writable_chunk = writable_chunk.compute().chunk()