Adding type suppressions for pyrefly

airio/_src/core/data_sources.py

@@ -22,7 +22,7 @@
 class DataSource(Protocol):
   """Interface for data sources wrappers with multiple splits support."""
 
-  splits: Iterable[str] = None
+  splits: Iterable[str] = None  # pyrefly: ignore[bad-assignment]
 
   def get_data_source(self, split: str):
     ...

airio/_src/core/dataset_iterators.py

@@ -24,7 +24,7 @@
 class AirIODatasetIterator(clu_dataset_iterator.DatasetIterator):
   """Wrapper iterator for AirIO."""
 
-  _iterator: collections.abc.Iterator[Any] = None
+  _iterator: collections.abc.Iterator[Any] = None  # pyrefly: ignore[bad-assignment]
 
   def __next__(self) -> clu_dataset_iterator.Element:
     raise NotImplementedError()

airio/_src/core/dataset_providers.py

@@ -43,12 +43,12 @@ class ShardInfo:
 class DatasetProviderBase(Protocol):
   """Abstract base for classes that provide a dataset."""
 
-  splits: Iterable[str] = None
+  splits: Iterable[str] = None  # pyrefly: ignore[bad-assignment]
 
   def get_dataset(
       self,
       sequence_lengths: Mapping[str, int] | None = None,
-      split: str = tfds.Split.TRAIN,
+      split: str = tfds.Split.TRAIN,  # pyrefly: ignore[missing-attribute]
       runtime_preprocessors: Sequence[grain.Transformation] | None = None,
       batch_size: int | None = None,
       shuffle: bool = True,
@@ -114,7 +114,7 @@ def __init__(
     all_tasks = [t for t in tasks if isinstance(t, Task)]
     all_mixtures = [m for m in tasks if isinstance(m, Mixture)]
     sub_tasks = [mix.leaf_tasks for mix in all_mixtures]
-    leaf_tasks = sum(sub_tasks, all_tasks)
+    leaf_tasks = sum(sub_tasks, all_tasks)  # pyrefly: ignore[no-matching-overload]
     duplicate_tasks = [
         t for t, c in collections.Counter(leaf_tasks).items() if c > 1
     ]
@@ -128,7 +128,7 @@ def __init__(
     self._proportions = dict(zip(tasks, proportions))
 
   def num_input_examples(self, split: str) -> int | None:
-    return sum(
+    return sum(  # pyrefly: ignore[no-matching-overload]
         t.num_input_examples(split)
         for t in self.tasks_or_mixtures
         if split in t.splits
@@ -164,14 +164,14 @@ def leaf_tasks(self) -> Sequence[Task]:
     tasks = [t for t in all_ if isinstance(t, Task)]
     mixtures = [m for m in all_ if isinstance(m, Mixture)]
     sub_tasks = [mix.leaf_tasks for mix in mixtures]
-    return sum(sub_tasks, tasks)
+    return sum(sub_tasks, tasks)  # pyrefly: ignore[no-matching-overload]
 
   @property
   def total_proportion(self) -> float:
     return sum(self._proportions.values())
 
   @property
-  def splits(self) -> Sequence[str]:
+  def splits(self) -> Sequence[str]:  # pyrefly: ignore[bad-override]
     splits = set()
     for task in self.tasks_or_mixtures:
       splits.update(task.splits)
@@ -218,7 +218,7 @@ def build(self) -> Task:
     if self._preprocessors is None:
       raise ValueError("Preprocessors have not been set on this task builder.")
 
-    return Task(
+    return Task(  # pyrefly: ignore[bad-instantiation]
         name=self._task_name,
         source=self._source,
         preprocessors=self._preprocessors,
@@ -248,7 +248,7 @@ def from_task(cls, task: Task) -> "TaskBuilder":
     Args:
       task: Existing task object.
     """
-    return TaskBuilder(
+    return TaskBuilder(  # pyrefly: ignore[bad-instantiation]
         task_name=task.name,
         source=task.source,
         preprocessors=task.get_preprocessors(),

airio/_src/core/preprocessors.py

@@ -60,7 +60,7 @@ def replace(self, **kwargs):
 
 @dataclasses.dataclass
 @typing.runtime_checkable
-class MapFnTransform(Protocol):
+class MapFnTransform(Protocol):  # pyrefly: ignore[bad-class-definition]
   """Transform to represent AirIO map preprocessors.
 
   Attrs:
@@ -84,7 +84,7 @@ def map(self, element):
 
 @dataclasses.dataclass
 @typing.runtime_checkable
-class RandomMapFnTransform(Protocol):
+class RandomMapFnTransform(Protocol):  # pyrefly: ignore[bad-class-definition]
   """Transform to represent AirIO random map preprocessors.
 
   Attrs:
@@ -108,7 +108,7 @@ def random_map(self, element, rng: np.random.Generator):
 
 @dataclasses.dataclass
 @typing.runtime_checkable
-class FilterFnTransform(Protocol):
+class FilterFnTransform(Protocol):  # pyrefly: ignore[bad-class-definition]
   """Transform to represent AirIO filter preprocessors.
 
   Attrs:

airio/_src/core/test_utils.py

@@ -35,7 +35,7 @@ def assert_datasets_equal(
   """
 
   if not isinstance(expected, list):
-    expected = [expected]
+    expected = [expected]  # pyrefly: ignore[bad-assignment]
   actual = list(dataset)
   absltest.TestCase().assertEqual(len(actual), len(expected))
 
@@ -79,4 +79,4 @@ def create_airio_injected_runtime_args(
       "batch_size": batch_size,
   }
   args = {k: provided[k] if provided[k] else defaults[k] for k in defaults}
-  return preprocessors.AirIOInjectedRuntimeArgs(**args)
+  return preprocessors.AirIOInjectedRuntimeArgs(**args)  # pyrefly: ignore[bad-argument-type]

airio/_src/core/tokenizer.py

@@ -38,7 +38,7 @@ def vocabulary(self) -> vocabularies.Vocabulary:
 
 @typing.runtime_checkable
 @dataclasses.dataclass(frozen=True)
-class Tokenizer(Generic[Inp, Out], Protocol):
+class Tokenizer(Generic[Inp, Out], Protocol):  # pyrefly: ignore[bad-class-definition]
   """Tokenizer class for AirIO tasks/mixtures."""
 
   tokenizer_configs: Mapping[str, TokenizerConfig]

airio/_src/core/vocabularies.py

@@ -237,7 +237,7 @@ def _load_model(
       return cls._ModelContext(tokenizer=tokenizer, sp_model=sp_model)
 
   @property
-  def pad_id(self) -> int | None:
+  def pad_id(self) -> int | None:  # pyrefly: ignore[bad-override]
     return PAD_ID
 
   @property
@@ -287,7 +287,7 @@ def __eq__(self, other):
     if not isinstance(other, SentencePieceVocabulary):
       return False
     try:
-      their_md5 = hashlib.md5(other.sp_model).hexdigest()
+      their_md5 = hashlib.md5(other.sp_model).hexdigest()  # pyrefly: ignore[bad-argument-type]
     # If other has no sp_model attribute, we can't test for equality
     except AttributeError:
       return False
@@ -298,6 +298,7 @@ def __eq__(self, other):
 
   def __str__(self) -> str:
     return (
+        # pyrefly: ignore[bad-argument-type]
         f"SentencePieceVocabulary(file={self.sentencepiece_model_file}, "
         f"extra_ids={self._extra_ids}, "
         f"spm_md5={hashlib.md5(self.sp_model).hexdigest()})"
@@ -344,7 +345,7 @@ def eos_id(self) -> int | None:
     return None
 
   @property
-  def pad_id(self) -> int | None:
+  def pad_id(self) -> int | None:  # pyrefly: ignore[bad-override]
     return PAD_ID
 
   @property

airio/_src/pygrain/common/feature_converters.py

@@ -101,22 +101,22 @@ def get_t5x_enc_dec_feature_converter_preprocessors(
     )
     pack_prep.append(
         preprocessors.LazyIterTransform(
-            packer, update_runtime_args=packer.update_runtime_args
+            packer, update_runtime_args=packer.update_runtime_args  # pyrefly: ignore[bad-argument-type]
         )
     )
   return (
       [
           preprocessors.MapFnTransform(
-              common_preprocessors.remove_features_not_in_sequence_lengths
+              common_preprocessors.remove_features_not_in_sequence_lengths  # pyrefly: ignore[bad-argument-count]
           ),
-          preprocessors.MapFnTransform(common_preprocessors.trim),
+          preprocessors.MapFnTransform(common_preprocessors.trim),  # pyrefly: ignore[bad-argument-count]
       ]
       + pack_prep
       + [
-          preprocessors.MapFnTransform(pad),
+          preprocessors.MapFnTransform(pad),  # pyrefly: ignore[bad-argument-count]
           preprocessors.MapFnTransform(
-              convert_features,
-              update_runtime_args=update_runtime_args,
+              convert_features,  # pyrefly: ignore[bad-argument-count]
+              update_runtime_args=update_runtime_args,  # pyrefly: ignore[unexpected-keyword]
           ),
       ]
   )
@@ -167,22 +167,22 @@ def get_t5x_lm_feature_converter_preprocessors(
     )
     packer_prep.append(
         preprocessors.LazyIterTransform(
-            packer, update_runtime_args=packer.update_runtime_args
+            packer, update_runtime_args=packer.update_runtime_args  # pyrefly: ignore[bad-argument-type]
         )
     )
   return (
       [
           preprocessors.MapFnTransform(
-              common_preprocessors.remove_features_not_in_sequence_lengths
+              common_preprocessors.remove_features_not_in_sequence_lengths  # pyrefly: ignore[bad-argument-count]
           ),
-          preprocessors.MapFnTransform(common_preprocessors.trim),
+          preprocessors.MapFnTransform(common_preprocessors.trim),  # pyrefly: ignore[bad-argument-count]
       ]
       + packer_prep
       + [
-          preprocessors.MapFnTransform(pad),
+          preprocessors.MapFnTransform(pad),  # pyrefly: ignore[bad-argument-count]
           preprocessors.MapFnTransform(
-              convert_features,
-              update_runtime_args=update_runtime_args,
+              convert_features,  # pyrefly: ignore[bad-argument-count]
+              update_runtime_args=update_runtime_args,  # pyrefly: ignore[unexpected-keyword]
           ),
       ]
   )
@@ -258,9 +258,9 @@ def swap_inputs_width(ex: dict[str, np.ndarray], old_val: int, new_val: int):
 
   preps = [
       preprocessors.MapFnTransform(
-          concat_and_add_masks, update_runtime_args=concat_task_feature_lengths
+          concat_and_add_masks, update_runtime_args=concat_task_feature_lengths  # pyrefly: ignore[bad-argument-count, unexpected-keyword]
       ),
-      preprocessors.MapFnTransform(replace_0s),
+      preprocessors.MapFnTransform(replace_0s),  # pyrefly: ignore[bad-argument-count]
   ]
   if pack:
     packer = (
@@ -269,15 +269,15 @@ def swap_inputs_width(ex: dict[str, np.ndarray], old_val: int, new_val: int):
         else packing.SingleBinTruePackIterPreprocessor
     )
     packer_prep = preprocessors.LazyIterTransform(
-        packer, update_runtime_args=packer.update_runtime_args
+        packer, update_runtime_args=packer.update_runtime_args  # pyrefly: ignore[bad-argument-type]
     )
-    preps.append(packer_prep)
+    preps.append(packer_prep)  # pyrefly: ignore[bad-argument-type]
   preps.extend([
-      preprocessors.MapFnTransform(common_preprocessors.trim),
-      preprocessors.MapFnTransform(pad),
-      preprocessors.MapFnTransform(restore_0s),
+      preprocessors.MapFnTransform(common_preprocessors.trim),  # pyrefly: ignore[bad-argument-count]
+      preprocessors.MapFnTransform(pad),  # pyrefly: ignore[bad-argument-count]
+      preprocessors.MapFnTransform(restore_0s),  # pyrefly: ignore[bad-argument-count]
       preprocessors.MapFnTransform(
-          convert_features, update_runtime_args=update_runtime_args
+          convert_features, update_runtime_args=update_runtime_args  # pyrefly: ignore[bad-argument-count, unexpected-keyword]
       ),
   ])
   return preps

airio/_src/pygrain/common/packing.py

@@ -221,7 +221,7 @@ def parent(self):
   def __len__(self):
     return len(self.parent)
 
-  def __getitem__(self, index: slice):
+  def __getitem__(self, index: slice):  # pyrefly: ignore[bad-override]
     if isinstance(index, slice):
       return self.slice(index)
     return self._packed_ds[index]
@@ -350,7 +350,7 @@ def get_state(self):
 
   def set_state(self, state):
     self._parent_iter.set_state(state["parent"])
-    self._packer = self._packer_type.from_dict(state["packer"])
+    self._packer = self._packer_type.from_dict(state["packer"])  # pyrefly: ignore[bad-argument-type]
     self._packed_examples = collections.deque[PyTree[np.ndarray]](
         [load_np_tree(t) for t in state["packed_examples"]]
     )
@@ -514,7 +514,7 @@ def fit_example(self, ex: PyTree[np.ndarray]) -> Sequence[PyTree[np.ndarray]]:
     # Add if example fits an existing partially packed example; check if
     # resulting partially packed example becomes fully packed
     fits = False
-    fully_packed: PartiallyPackedExample = None
+    fully_packed: PartiallyPackedExample = None  # pyrefly: ignore[bad-assignment]
     fully_packed_idx = None
     for idx, partially_packed in enumerate(self._partially_packed_examples):
       if partially_packed.example_fits(flat_ex):
@@ -532,7 +532,7 @@ def fit_example(self, ex: PyTree[np.ndarray]) -> Sequence[PyTree[np.ndarray]]:
           fully_packed.pack(),
           length_struct=self.feature_lengths,
       )
-      del self._partially_packed_examples[fully_packed_idx]
+      del self._partially_packed_examples[fully_packed_idx]  # pyrefly: ignore[unsupported-operation]
       # self._partially_packed_examples.remove(fully_packed)
       packed_examples.append(packed)
 
@@ -615,7 +615,7 @@ def __init__(
   ):
     self._feature_lengths = feature_lengths
     self._flat_feature_lengths = flatten(feature_lengths)
-    self._partially_packed_example: PartiallyPackedExample = None
+    self._partially_packed_example: PartiallyPackedExample = None  # pyrefly: ignore[bad-assignment]
     if feature_lengths:
       self._partially_packed_example = PartiallyPackedExample(
           copy.copy(self._flat_feature_lengths)

airio/_src/pygrain/data_sources.py

@@ -39,7 +39,7 @@ def __init__(
 
     self.splits = frozenset(self._split_to_filepattern.keys())
     self._sources = {
-        split: grain.ArrayRecordDataSource(self._split_to_filepattern[split])
+        split: grain.ArrayRecordDataSource(self._split_to_filepattern[split])  # pyrefly: ignore[bad-argument-type]
         for split in self.splits
     }
 
@@ -119,7 +119,7 @@ def __init__(
     self.splits = frozenset(self._split_to_filepattern.keys())
     self._sources = {}
     for split in self.splits:
-      json_data = json.load(Open(self._split_to_filepattern[split]))
+      json_data = json.load(Open(self._split_to_filepattern[split]))  # pyrefly: ignore[bad-argument-type]
       json_data = [json.dumps(d) for d in json_data]
       self._sources[split] = grain.InMemoryDataSource(elements=json_data)
 

airio/_src/pygrain/dataset_iterators.py

@@ -98,13 +98,13 @@ def peek_async(
 
   def get_state(self) -> Mapping[str, Any]:
     if self._state_as_dict:
-      return self._iterator.get_state()
-    return json.loads(self._iterator.get_state().decode())
+      return self._iterator.get_state()  # pyrefly: ignore[missing-attribute]
+    return json.loads(self._iterator.get_state().decode())  # pyrefly: ignore[missing-attribute]
 
   def set_state(self, state: Mapping[str, Any]) -> None:
     if not self._state_as_dict:
-      state = json.dumps(state, indent=4).encode()
-    self._iterator.set_state(state)
+      state = json.dumps(state, indent=4).encode()  # pyrefly: ignore[bad-assignment]
+    self._iterator.set_state(state)  # pyrefly: ignore[missing-attribute]
 
   def save(self, filename: epath.PathLike):
     filename = epath.Path(filename)