From ff4a5d3cc584a8a991fd5202ca8d87a4253b9c1c Mon Sep 17 00:00:00 2001 From: Oleh Prypin Date: Wed, 7 Oct 2026 04:14:53 -0700 Subject: [PATCH] Replace `# pytype: disable` suppressions with `# pyrefly: ignore` PiperOrigin-RevId: 995012451 --- .../python/dataset/transformations/mix_test.py | 18 +++++++----------- .../dataset/transformations/packing_test.py | 4 ++-- .../python/dataset/transformations/prefetch.py | 6 ++---- .../dataset/transformations/prefetch_test.py | 10 +++++----- .../transformations/process_prefetch.py | 14 ++++++-------- .../dataset/transformations/shuffle_test.py | 4 ++-- .../dataset/transformations/source_test.py | 6 ++---- .../python/dataset/transformations/zip_test.py | 14 +++++++------- 8 files changed, 33 insertions(+), 43 deletions(-) diff --git a/grain/_src/python/dataset/transformations/mix_test.py b/grain/_src/python/dataset/transformations/mix_test.py index a2530e6c0..f0c0b7dd4 100644 --- a/grain/_src/python/dataset/transformations/mix_test.py +++ b/grain/_src/python/dataset/transformations/mix_test.py @@ -165,13 +165,9 @@ def test_getitem_with_numpy_integer_index_does_not_overflow(self): ) for index_type in (np.int32, np.uint32, np.int64): index = index_type(3) - # pytype: disable=wrong-arg-types - # pytype: disable=unsupported-operands - input_index, mapped_index = selection_map[index] + input_index, mapped_index = selection_map[index] # pyrefly: ignore[bad-index] self.assertEqual(input_index, 1) self.assertEqual(mapped_index, 1) - # pytype: enable=unsupported-operands - # pytype: enable=wrong-arg-types class MixedMapDatasetTest(parameterized.TestCase): @@ -474,7 +470,7 @@ def test_getitem_with_numpy_integer_index_does_not_overflow(self): ) for dtype in [np.int32, np.uint32, np.int64]: index = dtype(5) - self.assertEqual(mixed_lzds[index], 5) # pytype: disable=unsupported-operands,wrong-arg-types + self.assertEqual(mixed_lzds[index], 5) class MixedIterDatasetTest(absltest.TestCase): @@ -608,11 +604,11 @@ def test_checkpointing(self): checkpoints = [] for _ in range(max_steps): - checkpoints.append(ds_iter.get_state()) # pytype: disable=attribute-error + checkpoints.append(ds_iter.get_state()) values_without_interruption.append(next(ds_iter)) for starting_step in [0, 1, 5, 8]: - ds_iter.set_state(checkpoints[starting_step]) # pytype: disable=attribute-error + ds_iter.set_state(checkpoints[starting_step]) for i in range(starting_step, max_steps): np.testing.assert_array_equal( next(ds_iter), values_without_interruption[i] @@ -766,13 +762,13 @@ def test_get_set_shard_states(self): self.assertEqual(next(it), 2) self.assertEqual(next(it), 12) - shard_states = it.get_shard_states() # pytype: disable=attribute-error + shard_states = it.get_shard_states() # pyrefly: ignore[missing-attribute] self.assertLen(shard_states, 2) self.assertLen(shard_states[0]["parent_states"], 2) self.assertLen(shard_states[1]["parent_states"], 2) it2 = mix_ds.__iter__() - it2.set_shard_states(shard_states) # pytype: disable=attribute-error + it2.set_shard_states(shard_states) # pyrefly: ignore[missing-attribute] self.assertEqual(next(it2), 4) self.assertEqual(next(it2), 14) @@ -853,7 +849,7 @@ def test_concat_selection_map_with_numpy_integer_index_does_not_overflow( # Prevents regression of xid/273443246 for dtype in [np.int32, np.uint32, np.int64]: index = dtype(3) - self.assertEqual(selection_map[index], (1, 1)) # pytype: disable=unsupported-operands,wrong-arg-types + self.assertEqual(selection_map[index], (1, 1)) # pyrefly: ignore[bad-index] if __name__ == "__main__": diff --git a/grain/_src/python/dataset/transformations/packing_test.py b/grain/_src/python/dataset/transformations/packing_test.py index 82aa2a527..38599ff83 100644 --- a/grain/_src/python/dataset/transformations/packing_test.py +++ b/grain/_src/python/dataset/transformations/packing_test.py @@ -49,7 +49,7 @@ def test_get_packed_batch_size_bytes(self): iterator = iterator._parent # 2*10*8 (values) + 2*10*4 (segment_ids) + 2*10*4 (positions) + 2*8 # (first_free_cell) = 160 + 80 + 80 + 16 = 336 - self.assertEqual(iterator.get_packed_batch_size_bytes(), 336) # pytype: disable=attribute-error + self.assertEqual(iterator.get_packed_batch_size_bytes(), 336) # pyrefly: ignore[missing-attribute] def test_get_packed_batch_size_bytes_before_next(self): ds = source.SourceMapDataset([ # pyrefly: ignore[bad-argument-type] @@ -66,7 +66,7 @@ def test_get_packed_batch_size_bytes_before_next(self): while not isinstance(iterator, packing.PackingDatasetIterator): iterator = iterator._parent # Check size before calling next() - self.assertRaises(ValueError, iterator.get_packed_batch_size_bytes) # pytype: disable=attribute-error + self.assertRaises(ValueError, iterator.get_packed_batch_size_bytes) # pyrefly: ignore[missing-attribute] class FirstFitPackIterDatasetTest( diff --git a/grain/_src/python/dataset/transformations/prefetch.py b/grain/_src/python/dataset/transformations/prefetch.py index 1b2efc78b..013e68a95 100644 --- a/grain/_src/python/dataset/transformations/prefetch.py +++ b/grain/_src/python/dataset/transformations/prefetch.py @@ -539,7 +539,6 @@ def __init__( variable_size_queue.VariableSizeQueue | queue.Queue[BufferElementT] ) = queue.Queue(maxsize=self._target_prefetch_buffer_size) - # pytype: disable=attribute-error # pylint: disable=protected-access def _initialize_stats( @@ -557,10 +556,9 @@ def _initialize_stats( @functools.cached_property def _stats(self): return self._initialize_stats( - self._ctx.dataset_options.execution_tracking_mode + self._ctx.dataset_options.execution_tracking_mode # pyrefly: ignore[bad-argument-type] ) - # pytype: enable=attribute-error # pylint: enable=protected-access def start_prefetch(self): @@ -653,7 +651,7 @@ def _stop_prefetch(self, clear_buffer: bool = True): else: assert isinstance(self._buffer, variable_size_queue.VariableSizeQueue) # Increase the buffer size by 1 to unblock the producer. - self._buffer.set_max_size(self._target_prefetch_buffer_size + 1) # pytype: disable=attribute-error + self._buffer.set_max_size(self._target_prefetch_buffer_size + 1) if not sys.is_finalizing(): # Joining the worker thread is not necessary when the Python interpreter diff --git a/grain/_src/python/dataset/transformations/prefetch_test.py b/grain/_src/python/dataset/transformations/prefetch_test.py index 328d2ff20..168d5871b 100644 --- a/grain/_src/python/dataset/transformations/prefetch_test.py +++ b/grain/_src/python/dataset/transformations/prefetch_test.py @@ -241,11 +241,11 @@ def test_checkpoint(self): values_without_interruption = [] checkpoints = [] for _ in range(max_steps): - checkpoints.append(ds_iter.get_state()) # pytype: disable=attribute-error + checkpoints.append(ds_iter.get_state()) values_without_interruption.append(next(ds_iter)) for starting_step in [0, 1, 5, 12, 18]: - ds_iter.set_state(checkpoints[starting_step]) # pytype: disable=attribute-error + ds_iter.set_state(checkpoints[starting_step]) for i in range(starting_step, max_steps): value = next(ds_iter) self.assertEqual(value, values_without_interruption[i]) @@ -258,7 +258,7 @@ def test_invalid_checkpoint(self, next_index: int): f'Checkpoint `next_index` {next_index} is out of range for dataset of' ' length 20.', ): - ds_iter.set_state({'next_index': next_index}) # pytype: disable=attribute-error + ds_iter.set_state({'next_index': next_index}) def test_filter_all_elements_warns(self): ds = ( @@ -405,7 +405,7 @@ def test_start_prefetch( it.start_prefetch() # Check that the buffer was filled before we start processing elements. if num_threads > 0: - self.assertNotEmpty(it._buffer) # pytype: disable=attribute-error + self.assertNotEmpty(it._buffer) # pyrefly: ignore[missing-attribute] self.assertEqual(list(it), list(range(10))) @@ -677,7 +677,7 @@ def test_no_mem_leak_with_double_prefetch(self, close: bool): for _ in range(5): _ = next(it) if close: - it.close() # pytype: disable=attribute-error + it.close() @absltest.skipIf(platform.system() == 'Darwin', 'Fails on macos-14 runner.') @parameterized.parameters([True, False]) diff --git a/grain/_src/python/dataset/transformations/process_prefetch.py b/grain/_src/python/dataset/transformations/process_prefetch.py index 1ae1ce388..93e0f44e9 100644 --- a/grain/_src/python/dataset/transformations/process_prefetch.py +++ b/grain/_src/python/dataset/transformations/process_prefetch.py @@ -246,7 +246,7 @@ def _put_dataset_elements_in_buffer( new_state_or_index = set_state_queue.get() parent_exhausted = False if new_state_or_index is not None: - if not grain_queue.add_element_to_queue( # pytype: disable=wrong-arg-types + if not grain_queue.add_element_to_queue( (_SetStateIsDone(), None, None, None, None), buffer, # pyrefly: ignore[bad-argument-type] should_stop.is_set, @@ -266,7 +266,7 @@ def _put_dataset_elements_in_buffer( try: element = next(it) except Exception as e: # pylint: disable=broad-except - grain_queue.add_element_to_queue( # pytype: disable=wrong-arg-types + grain_queue.add_element_to_queue( ( None, None, @@ -289,7 +289,7 @@ def _put_dataset_elements_in_buffer( it._stats.record_bytes_produced(element) # pylint: disable=protected-access if next_index is not None: next_index += 1 - if not grain_queue.add_element_to_queue( # pytype: disable=wrong-arg-types + if not grain_queue.add_element_to_queue( (element, it.get_state(), next_index, None, None), buffer, # pyrefly: ignore[bad-argument-type] should_stop.is_set, @@ -301,7 +301,7 @@ def _put_dataset_elements_in_buffer( except Exception as e: # pylint: disable=broad-except _clear_queue_and_maybe_unlink_shm(buffer) _clear_queue_and_maybe_unlink_shm(set_state_queue) - grain_queue.add_element_to_queue( # pytype: disable=wrong-arg-types + grain_queue.add_element_to_queue( ( None, None, @@ -401,7 +401,6 @@ def __init__( def buffer_size(self) -> int: return self._buffer_size - # pytype: disable=attribute-error # pylint: disable=protected-access def _initialize_stats( self, execution_tracking_mode: base.ExecutionTrackingMode @@ -431,10 +430,9 @@ def _initialize_stats( @functools.cached_property def _stats(self): return self._initialize_stats( - self._ctx.dataset_options.execution_tracking_mode + self._ctx.dataset_options.execution_tracking_mode # pyrefly: ignore[bad-argument-type] ) - # pytype: enable=attribute-error # pylint: disable=protected-access def start_prefetch(self) -> None: @@ -675,7 +673,7 @@ def __init__( def __iter__(self) -> dataset.DatasetIterator[T]: if not _is_in_worker_process: return self._parent.__iter__() - dataset.set_slice(self._parent, self._slice, self._sequential_slice) # pyrefly: ignore[bad-argument-type] + dataset.set_slice(self._parent, self._slice, self._sequential_slice) return self._parent.__iter__() @property diff --git a/grain/_src/python/dataset/transformations/shuffle_test.py b/grain/_src/python/dataset/transformations/shuffle_test.py index 9e578ee5c..6ba1ef3a0 100644 --- a/grain/_src/python/dataset/transformations/shuffle_test.py +++ b/grain/_src/python/dataset/transformations/shuffle_test.py @@ -81,7 +81,7 @@ def test_shuffled_index_with_numpy_integer_does_not_overflow(self): for dtype in [np.int32, np.uint32, np.int64]: index = dtype(5) # Should not raise OverflowError when computing modulo 2**32 - _ = ds._shuffled_index(index) # pytype: disable=wrong-arg-types + _ = ds._shuffled_index(index) # pyrefly: ignore[bad-argument-type] def test_cross_version_determinism(self): # This test validates shuffle determinism across different versions of @@ -250,7 +250,7 @@ def test_shuffled_index_with_numpy_integer_does_not_overflow(self): # Prevents regression of xid/273443246 for dtype in [np.int32, np.uint32, np.int64]: index = dtype(5) - _ = ds._shuffled_index(index) # pytype: disable=wrong-arg-types + _ = ds._shuffled_index(index) # pyrefly: ignore[bad-argument-type] class WindowShuffleInterDatasetTest(parameterized.TestCase): diff --git a/grain/_src/python/dataset/transformations/source_test.py b/grain/_src/python/dataset/transformations/source_test.py index 88bd4189a..0d2623cfb 100644 --- a/grain/_src/python/dataset/transformations/source_test.py +++ b/grain/_src/python/dataset/transformations/source_test.py @@ -59,10 +59,8 @@ def setUp(self): self._batched_read_data_source = ( _MySupportsBatchedReadRandomAccessDataSource([1, 2, 3, 4, 5]) ) - self._lazy_dataset_source = source.SourceMapDataset( # pytype: disable=wrong-arg-types - self._data_source - ) - self._batched_read_lazy_dataset_source = source.SourceMapDataset( # pytype: disable=wrong-arg-types + self._lazy_dataset_source = source.SourceMapDataset(self._data_source) + self._batched_read_lazy_dataset_source = source.SourceMapDataset( self._batched_read_data_source ) diff --git a/grain/_src/python/dataset/transformations/zip_test.py b/grain/_src/python/dataset/transformations/zip_test.py index b0e3e8426..8f9a702a0 100644 --- a/grain/_src/python/dataset/transformations/zip_test.py +++ b/grain/_src/python/dataset/transformations/zip_test.py @@ -296,11 +296,11 @@ def test_get_set_shard_states(self): self.assertEqual(next(it), (0, 10)) self.assertEqual(next(it), (2, 12)) - shard_states = it.get_shard_states() # pytype: disable=attribute-error + shard_states = it.get_shard_states() # pyrefly: ignore[missing-attribute] self.assertLen(shard_states, 2) it2 = zipped_ds.__iter__() - it2.set_shard_states(shard_states) # pytype: disable=attribute-error + it2.set_shard_states(shard_states) # pyrefly: ignore[missing-attribute] self.assertEqual(next(it2), (4, 14)) def test_get_set_shard_states_nested(self): @@ -348,11 +348,11 @@ def set_state(self, state): self.assertEqual(next(it), (0, 10)) self.assertEqual(next(it), (2, 12)) - shard_states = it.get_shard_states() # pytype: disable=attribute-error + shard_states = it.get_shard_states() # pyrefly: ignore[missing-attribute] self.assertLen(shard_states, 2) it2 = zipped_ds.__iter__() - it2.set_shard_states(shard_states) # pytype: disable=attribute-error + it2.set_shard_states(shard_states) # pyrefly: ignore[missing-attribute] self.assertEqual(next(it2), (4, 14)) def test_empty_parents(self): @@ -385,7 +385,7 @@ def test_get_shard_states_unsupported_parent(self): with self.assertRaisesRegex( ValueError, "does not support elastic resizing" ): - it.get_shard_states() # pytype: disable=attribute-error + it.get_shard_states() # pyrefly: ignore[missing-attribute] def test_get_shard_states_mismatched_shard_counts(self): class DummyElasticDataset(dataset.IterDataset): @@ -425,7 +425,7 @@ def set_shard_states(self, shard_states): with self.assertRaisesWithLiteralMatch( ValueError, "All parents must have the same number of shards." ): - it.get_shard_states() # pytype: disable=attribute-error + it.get_shard_states() # pyrefly: ignore[missing-attribute] def test_set_shard_states_empty(self): ds1 = [ @@ -434,7 +434,7 @@ def test_set_shard_states_empty(self): interleave_ds1 = interleave.InterleaveIterDataset(ds1, cycle_length=2) zipped_ds = zip_ds.ZipIterDataset([interleave_ds1, interleave_ds1]) it = iter(zipped_ds) - it.set_shard_states([]) # pytype: disable=attribute-error + it.set_shard_states([]) # pyrefly: ignore[missing-attribute] def test_str(self): ds = zip_ds.ZipIterDataset(