Replace # pytype: disable suppressions with # pyrefly: ignore

tensorflow_datasets/core/community/dataset_sources.py

@@ -55,7 +55,7 @@ def to_json(self) -> utils.JsonValue:
     if len(self.filenames) == 1:
       return os.fspath(self.root_path / self.filenames[0])
     else:
-      return {  # pyrefly: ignore[bad-return]
+      return {
           'root_path': os.fspath(self.root_path),
           'filenames': self.filenames,  # pyrefly: ignore[bad-assignment]
       }

tensorflow_datasets/core/community/register_package.py

@@ -187,7 +187,7 @@ def refresh(self) -> None:
       return
 
     # If read was sucessful, update the cache with the new dataset list
-    self._cached_path.write_text(content)  # pytype: disable=attribute-error  # attribute-variable-annotations
+    self._cached_path.write_text(content)
     self._refresh_from_content(content)
 
 
@@ -260,7 +260,7 @@ def builder(
       **builder_kwargs: Any,
   ) -> dataset_builder.DatasetBuilder:
     """Returns the dataset builder."""
-    return self.builder_cls(name)(**builder_kwargs)  # pytype: disable=not-instantiable
+    return self.builder_cls(name)(**builder_kwargs)
 
 
 def list_ds_packages_for_namespace(

tensorflow_datasets/core/dataset_builders/conll/conll_dataset_builder.py

@@ -97,7 +97,7 @@ class ConllDatasetBuilder(
   @property
   def builder_config(self) -> ConllBuilderConfig:
     """`tfds.core.BuilderConfig` for this builder."""
-    return self._builder_config  # pytype: disable=bad-return-type  # always-use-return-annotations
+    return self._builder_config  # pyrefly: ignore[bad-return]
 
   def create_dataset_info(
       self,

tensorflow_datasets/core/dataset_builders/conll/conllu_dataset_builder.py

@@ -158,7 +158,7 @@ class ConllUDatasetBuilder(
   @property
   def builder_config(self) -> ConllUBuilderConfig:
     """`tfds.core.BuilderConfig` for this builder."""
-    return self._builder_config  # pytype: disable=bad-return-type  # always-use-return-annotations
+    return self._builder_config  # pyrefly: ignore[bad-return]
 
   def create_dataset_info(
       self,

tensorflow_datasets/rlds/datasets/locomotion/locomotion.py

@@ -70,7 +70,7 @@ class Locomotion(tfds.core.GeneratorBasedBuilder):
           description=(
               'Dataset generated by a SAC agent trained for 1M steps for Ant.'
           ),
-          supervised_keys=None,  # pytype: disable=wrong-arg-types  # gen-stub-imports
+          supervised_keys=None,
       ),
       rlds_base.DatasetConfig(
           name='hopper_sac_1M_single_policy_stochastic',
@@ -85,7 +85,7 @@ class Locomotion(tfds.core.GeneratorBasedBuilder):
               'Dataset generated by a SAC agent trained for 1M steps for'
               ' Hopper.'
           ),
-          supervised_keys=None,  # pytype: disable=wrong-arg-types  # gen-stub-imports
+          supervised_keys=None,
       ),
       rlds_base.DatasetConfig(
           name='halfcheetah_sac_1M_single_policy_stochastic',
@@ -100,7 +100,7 @@ class Locomotion(tfds.core.GeneratorBasedBuilder):
               'Dataset generated by a SAC agent trained for 1M steps for'
               ' HalfCheetah.'
           ),
-          supervised_keys=None,  # pytype: disable=wrong-arg-types  # gen-stub-imports
+          supervised_keys=None,
       ),
       rlds_base.DatasetConfig(
           name='walker2d_sac_1M_single_policy_stochastic',
@@ -115,7 +115,7 @@ class Locomotion(tfds.core.GeneratorBasedBuilder):
               'Dataset generated by a SAC agent trained for 1M steps for'
               ' Walker2d.'
           ),
-          supervised_keys=None,  # pytype: disable=wrong-arg-types  # gen-stub-imports
+          supervised_keys=None,
       ),
       rlds_base.DatasetConfig(
           name='humanoid_sac_15M_single_policy_stochastic',
@@ -130,7 +130,7 @@ class Locomotion(tfds.core.GeneratorBasedBuilder):
               'Dataset generated by a SAC agent trained for 15M steps for'
               ' Humanoid.'
           ),
-          supervised_keys=None,  # pytype: disable=wrong-arg-types  # gen-stub-imports
+          supervised_keys=None,
       ),
   ]
 

tensorflow_datasets/rlds/datasets/robosuite_panda_pick_place_can/robosuite_panda_pick_place_can.py

@@ -72,7 +72,6 @@ class RobosuitePandaPickPlaceCan(tfds.core.GeneratorBasedBuilder):
       'synthetic_stochastic_sac_afe13968': 'https://storage.googleapis.com/rlds_external_data_release/rlds_robosuite_panda_pick_place_can_synthetic_stochastic_sac_afe13968.tar.gz',
   }
 
-  # pytype: disable=wrong-keyword-args
   BUILDER_CONFIGS = [
       rlds_base.DatasetConfig(
           name='human_dc29b40a',
@@ -114,22 +113,22 @@ class RobosuitePandaPickPlaceCan(tfds.core.GeneratorBasedBuilder):
               ),
           },
           action_info=tfds.features.Tensor(shape=(7,), dtype=np.float64),
-          reward_info=np.float64,
-          discount_info=np.float64,
+          reward_info=np.float64,  # pyrefly: ignore[bad-argument-type]
+          discount_info=np.float64,  # pyrefly: ignore[bad-argument-type]
           episode_metadata_info={
-              'agent_id': np.str_,
-              'episode_index': np.int32,
-              'episode_id': np.str_,
+              'agent_id': np.str_,  # pyrefly: ignore[bad-assignment]
+              'episode_index': np.int32,  # pyrefly: ignore[bad-assignment]
+              'episode_id': np.str_,  # pyrefly: ignore[bad-assignment]
           },
           step_metadata_info={
-              'tag:placed': np.bool_,
+              'tag:placed': np.bool_,  # pyrefly: ignore[bad-assignment]
               'image': tfds.features.Image(),
           },
           citation=_CITATION,
           homepage=_HOMEPAGE,
           overall_description=_DESCRIPTION,
           description='Human generated dataset (50 episodes).',
-          supervised_keys=None,  # pytype: disable=wrong-arg-types  # gen-stub-imports
+          supervised_keys=None,
       ),
       rlds_base.DatasetConfig(
           name='human_images_dc29b40a',
@@ -183,15 +182,15 @@ class RobosuitePandaPickPlaceCan(tfds.core.GeneratorBasedBuilder):
               ),
           },
           action_info=tfds.features.Tensor(shape=(7,), dtype=np.float64),
-          reward_info=np.float64,
-          discount_info=np.float64,
+          reward_info=np.float64,  # pyrefly: ignore[bad-argument-type]
+          discount_info=np.float64,  # pyrefly: ignore[bad-argument-type]
           episode_metadata_info={
-              'agent_id': np.str_,
-              'episode_index': np.int32,
-              'episode_id': np.str_,
+              'agent_id': np.str_,  # pyrefly: ignore[bad-assignment]
+              'episode_index': np.int32,  # pyrefly: ignore[bad-assignment]
+              'episode_id': np.str_,  # pyrefly: ignore[bad-assignment]
           },
           step_metadata_info={
-              'tag:placed': np.bool_,
+              'tag:placed': np.bool_,  # pyrefly: ignore[bad-assignment]
               'image': tfds.features.Image(),
           },
           citation=_CITATION,
@@ -201,7 +200,7 @@ class RobosuitePandaPickPlaceCan(tfds.core.GeneratorBasedBuilder):
               ' angles in the observation.'
               ' Note that it may take some time to generate.'
           ),
-          supervised_keys=None,  # pytype: disable=wrong-arg-types  # gen-stub-imports
+          supervised_keys=None,
       ),
       rlds_base.DatasetConfig(
           name='synthetic_stochastic_sac_afe13968',
@@ -243,15 +242,15 @@ class RobosuitePandaPickPlaceCan(tfds.core.GeneratorBasedBuilder):
               ),
           },
           action_info=tfds.features.Tensor(shape=(7,), dtype=np.float32),
-          reward_info=np.float64,
-          discount_info=np.float64,
+          reward_info=np.float64,  # pyrefly: ignore[bad-argument-type]
+          discount_info=np.float64,  # pyrefly: ignore[bad-argument-type]
           episode_metadata_info={
-              'agent_id': np.str_,
-              'episode_index': np.int32,
-              'episode_id': np.str_,
+              'agent_id': np.str_,  # pyrefly: ignore[bad-assignment]
+              'episode_index': np.int32,  # pyrefly: ignore[bad-assignment]
+              'episode_id': np.str_,  # pyrefly: ignore[bad-assignment]
           },
           step_metadata_info={
-              'tag:placed': np.bool_,
+              'tag:placed': np.bool_,  # pyrefly: ignore[bad-assignment]
               'image': tfds.features.Image(),
           },
           citation=_CITATION,
@@ -261,12 +260,10 @@ class RobosuitePandaPickPlaceCan(tfds.core.GeneratorBasedBuilder):
               'Synthetic dataset generated by a stochastic agent trained with'
               ' SAC (200 episodes).'
           ),
-          supervised_keys=None,  # pytype: disable=wrong-arg-types  # gen-stub-imports
+          supervised_keys=None,
       ),
   ]
 
-  # pytype: enable=wrong-keyword-args
-
 
   def _info(self) -> tfds.core.DatasetInfo:
     """Returns the dataset metadata."""

tensorflow_datasets/robomimic/dataset_utils.py

@@ -112,18 +112,16 @@ def make_builder_configs(dataset: DataSource):
   for task, details in TASKS.items():
     if dataset in details['datasets']:  # pyrefly: ignore[not-iterable]
       for observation_type in [ObservationType.IMAGE, ObservationType.LOW_DIM]:
-        # pytype: disable=wrong-keyword-args
         configs.append(
             # name is inherited from base dataclass unbeknownst to the linter
             BuilderConfig(  # pylint: disable=unexpected-keyword-arg
                 name=f'{task}_{dataset}_{observation_type}',
                 task=task,
                 dataset=dataset,
                 filename=observation_type,
-                horizon=details['horizon'],
+                horizon=details['horizon'],  # pyrefly: ignore[bad-argument-type]
             )
         )
-        # pytype: enable=wrong-keyword-args
   return configs
 
 
@@ -301,7 +299,7 @@ def _get_features(self) -> tfds.features.FeaturesDict:
     features = tfds.features.FeaturesDict({
         'horizon': np.int32,
         'episode_id': np.str_,
-        'steps': tfds.features.Dataset({  # pyrefly: ignore[bad-argument-type]
+        'steps': tfds.features.Dataset({
             'action': tensor_feature(action_size),  # pyrefly: ignore[bad-argument-type]
             'observation': observation,  # pyrefly: ignore[bad-assignment]
             'reward': np.float64,

tensorflow_datasets/scripts/deployment/export_community_datasets.py

@@ -36,7 +36,7 @@
 
 
 def main(_):
-  export_community_datasets(in_path=_IN_PATH, out_path=_OUT_PATH)  # pytype: disable=wrong-arg-types
+  export_community_datasets(in_path=_IN_PATH, out_path=_OUT_PATH)
 
 
 def export_community_datasets(

tensorflow_datasets/scripts/documentation/document_datasets.py

@@ -235,7 +235,7 @@ def get_config_builder(config) -> tfds.core.DatasetBuilder:
     return BuilderToDocument(
         sections=sections,
         namespace=None,
-        builder=builder_cls(),  # pytype: disable=not-instantiable
+        builder=builder_cls(),
         config_builders=[],
     )
 
@@ -365,30 +365,28 @@ def iter_documentation_builders(
   print('Retrieving the list of builders...')
   datasets = datasets or _all_tfds_datasets()
 
-  # pytype: disable=attribute-error
-  if doc_util_paths.fig_base_path:
+  if doc_util_paths.fig_base_path:  # pyrefly: ignore[missing-attribute]
     visu_doc_util = doc_utils.VisualizationDocUtil(
         base_path=doc_util_paths.fig_base_path,
-        base_url=doc_util_paths.fig_base_url,
+        base_url=doc_util_paths.fig_base_url,  # pyrefly: ignore[missing-attribute]
     )
   else:
     visu_doc_util = None
 
-  if doc_util_paths.df_base_path:
+  if doc_util_paths.df_base_path:  # pyrefly: ignore[missing-attribute]
     df_doc_util = doc_utils.DataframeDocUtil(
         base_path=doc_util_paths.df_base_path,
-        base_url=doc_util_paths.df_base_url,
+        base_url=doc_util_paths.df_base_url,  # pyrefly: ignore[missing-attribute]
     )
   else:
     df_doc_util = None
 
   if doc_util_paths.fig_base_path:
     nightly_doc_util = doc_utils.NightlyDocUtil(
-        path=doc_util_paths.nightly_path,
+        path=doc_util_paths.nightly_path,  # pyrefly: ignore[missing-attribute]
     )
   else:
     nightly_doc_util = None
-  # pytype: enable=attribute-error
 
   document_single_builder_fn = functools.partial(
       _document_single_builder,

tensorflow_datasets/structured/cherry_blossoms/cherry_blossoms.py

@@ -90,7 +90,7 @@ def _split_generators(self, dl_manager: tfds.download.DownloadManager):
 
   def _generate_examples(self, path):
     """Yields examples."""
-    with path.open() as f:  # pytype: disable=attribute-error  # gen-stub-imports
+    with path.open() as f:
       reader = csv.DictReader(f, delimiter=';')
       for index, row in enumerate(reader):
         example = {k: v.replace('NA', 'nan') for k, v in row.items()}

tensorflow_datasets/structured/howell/howell.py

@@ -89,7 +89,7 @@ def _split_generators(self, dl_manager: tfds.download.DownloadManager):
 
   def _generate_examples(self, file_path):
     """Yields examples."""
-    with file_path.open() as f:  # pytype: disable=attribute-error  # gen-stub-imports
+    with file_path.open() as f:
       reader = csv.DictReader(f, delimiter=";")
       for index, row in enumerate(reader):
         example = dict(row.items())

tensorflow_datasets/structured/wine_quality/wine_quality.py

@@ -119,7 +119,7 @@ def _info(self):
     return tfds.core.DatasetInfo(
         builder=self,
         description=_DESCRIPTION,
-        features=tfds.features.FeaturesDict({  # pyrefly: ignore[bad-argument-type]
+        features=tfds.features.FeaturesDict({
             "quality": np.int32,
             "features": features_dict,  # pyrefly: ignore[bad-assignment]
         }),

tensorflow_datasets/testing/huggingface/dummy_dataset.py

@@ -17,7 +17,7 @@
 
 from pathlib import Path  # < Direct import of path should be patched  # pylint: disable=g-importing-member
 
-import datasets  # pytype: disable=import-error
+import datasets  # pyrefly: ignore[missing-import]
 import tensorflow_datasets.public_api as tfds
 
 

tensorflow_datasets/text/c4_wsrs/c4_wsrs_utils.py

@@ -265,7 +265,7 @@ def extract_snippets(
         continue
       count_by_pair[('', '')] += 1
 
-    yield key, (snippet, rarest_pair, index_to_pair)  # pytype: disable=bad-return-type  # container-simplification
+    yield key, (snippet, rarest_pair, index_to_pair)  # pyrefly: ignore[invalid-yield]
 
 
 def _reverse_substitute_snippet(

tensorflow_datasets/text/dolphin_number_word/dolphin_number_word.py

@@ -90,7 +90,7 @@ def _split_generators(self, dl_manager: tfds.download.DownloadManager):
   def _generate_examples(self, json_paths, split):
     """Parses available JSON files in a given split, returns TF examples."""
     # Currently assume all files in archive are json-formatted.
-    for file in json_paths.glob('*' + split + '*'):  # pytype: disable=attribute-error  # gen-stub-imports
+    for file in json_paths.glob('*' + split + '*'):
       json_string = file.read_text()
       examples = json.loads(json_string)
       for e in examples: