Replace # pytype: disable suppressions with # pyrefly: ignore

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__":

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(

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

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])

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

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):

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
     )
 

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(