Automated Code Change

PiperOrigin-RevId: 944182406

Author: h-joo

tensorflow_datasets/core/as_dataframe.py

@@ -127,7 +127,7 @@ def _make_columns(
 ) -> List[ColumnInfo]:
   """Extract the columns info of the `panda.DataFrame`."""
   return [
-      ColumnInfo.from_spec(path, ds_info)
+      ColumnInfo.from_spec(path, ds_info)  # pyrefly: ignore[bad-argument-type]
       for path, _ in py_utils.flatten_with_path(specs)
   ]
 
@@ -210,7 +210,7 @@ def _repr_html_(self) -> Union[None, str]:
 
   # Flatten the keys names, specs,... while keeping the feature key definition
   # order
-  columns = _make_columns(ds.element_spec, ds_info=ds_info)
+  columns = _make_columns(ds.element_spec, ds_info=ds_info)  # pyrefly: ignore[bad-argument-type]
   rows = [_make_row_dict(ex, columns) for ex in dataset_utils.as_numpy(ds)]
   df = StyledDataFrame(rows)
   df.current_style.format({c.name: c.format_fn for c in columns if c.format_fn})

tensorflow_datasets/core/dataset_builder.py

@@ -136,8 +136,8 @@ def from_dataset_info(
         name=info_proto.config_name,
         description=info_proto.config_description,
         version=info_proto.version,
-        release_notes=info_proto.release_notes or {},
-        tags=info_proto.config_tags or [],
+        release_notes=info_proto.release_notes or {},  # pyrefly: ignore[bad-argument-type]
+        tags=info_proto.config_tags or [],  # pyrefly: ignore[bad-argument-type]
     )
 
 
@@ -292,7 +292,7 @@ def __init__(
     # Extract code version (VERSION or config)
     self._version = self._pick_version(version)
     # Compute the base directory (for download) and dataset/version directory.
-    self._data_dir_root, self._data_dir = self._build_data_dir(data_dir)
+    self._data_dir_root, self._data_dir = self._build_data_dir(data_dir)  # pyrefly: ignore[bad-argument-type]
     # If the dataset info is available, use it.
     dataset_info_path = dataset_info.dataset_info_path(self.data_path)
     if retry.retry(dataset_info_path.exists):
@@ -924,14 +924,14 @@ def build_single_data_source(split: str) -> Sequence[Any]:
         case file_adapters.FileFormat.ARRAY_RECORD:
           return array_record.ArrayRecordDataSource(
               info,
-              split=split,
+              split=split,  # pyrefly: ignore[bad-argument-type]
               decoders=decoders,
               deserialize_method=deserialize_method,
           )
         case file_adapters.FileFormat.PARQUET:
           return parquet.ParquetDataSource(
               info,
-              split=split,
+              split=split,  # pyrefly: ignore[bad-argument-type]
               decoders=decoders,
               deserialize_method=deserialize_method,
           )
@@ -1298,7 +1298,7 @@ def _make_download_manager(
       # Note: Error will be raised here if user try to record checksums
       # from a `zipapp`
       try:
-        register_checksums_path = utils.to_write_path(self._checksums_path)
+        register_checksums_path = utils.to_write_path(self._checksums_path)  # pyrefly: ignore[bad-argument-type]
         download.validate_checksums_path(register_checksums_path)
       except Exception:  # pylint: disable=broad-except
         raise
@@ -1558,7 +1558,7 @@ def _as_dataset(
 
     reader = reader_lib.Reader(
         example_specs=example_specs,
-        file_format=file_format,
+        file_format=file_format,  # pyrefly: ignore[bad-argument-type]
     )
     decode_fn = functools.partial(features.decode_example, decoders=decoders)
     return reader.read(
@@ -1849,7 +1849,7 @@ def read_tfrecord_as_dataset(
 
     return tf.data.TFRecordDataset(
         filenames=filenames,
-        compression_type=compression_type,
+        compression_type=compression_type,  # pyrefly: ignore[bad-argument-type]
         num_parallel_reads=num_parallel_reads,
     )
 
@@ -1939,7 +1939,7 @@ def _download_and_prepare(
         beam_runner=download_config.beam_runner,
         example_writer=self._example_writer(),
         # The following options are ignored by `ShardBasedBuilder`.
-        ignore_duplicates=None,
+        ignore_duplicates=None,  # pyrefly: ignore[bad-argument-type]
         max_examples_per_split=None,
         shard_config=None,
     )
@@ -2000,7 +2000,7 @@ class BeamBasedBuilder(GeneratorBasedBuilder):
   def _generate_examples(
       self, *args: Any, **kwargs: Any
   ) -> split_builder_lib.SplitGenerator:
-    return self._build_pcollection(*args, **kwargs)
+    return self._build_pcollection(*args, **kwargs)  # pyrefly: ignore[missing-attribute]
 
 
 def _check_split_names(split_names: Iterable[str]) -> None:

tensorflow_datasets/core/dataset_collection_builder.py

@@ -107,7 +107,7 @@ def from_cls(
     return cls(
         name=name,
         release_notes=release_notes,
-        description=description,
+        description=description,  # pyrefly: ignore[bad-argument-type]
         citation=citation,
         homepage=homepage,
     )

tensorflow_datasets/core/dataset_info.py

@@ -162,7 +162,7 @@ def from_proto(
         module_name=info_proto.module_name,
         config_name=info_proto.config_name,
         config_description=info_proto.config_description,
-        config_tags=info_proto.config_tags or [],
+        config_tags=info_proto.config_tags or [],  # pyrefly: ignore[bad-argument-type]
         release_notes={k: v for k, v in info_proto.release_notes.items()},
     )
 
@@ -728,7 +728,7 @@ def read_from_directory(self, dataset_info_dir: epath.PathLike) -> None:
     # Update splits
     filename_template = naming.ShardedFileTemplate(  # pytype: disable=wrong-arg-types  # always-use-property-annotation
         dataset_name=self.name,
-        data_dir=self.data_dir,
+        data_dir=self.data_dir,  # pyrefly: ignore[bad-argument-type]
         filetype_suffix=parsed_proto.file_format or "tfrecord",
     )
     split_dict = splits_lib.SplitDict.from_proto(
@@ -743,7 +743,7 @@ def read_from_directory(self, dataset_info_dir: epath.PathLike) -> None:
     # For `ReadOnlyBuilder`, reconstruct the features from the config.
     elif feature_lib.make_config_path(dataset_info_dir).exists():
       self._features = top_level_feature.TopLevelFeature.from_config(
-          dataset_info_dir
+          dataset_info_dir  # pyrefly: ignore[bad-argument-type]
       )
 
     # If the dataset was loaded from file, self.metadata will be `None`, so
@@ -827,7 +827,7 @@ def add_file_data_source_access(
     access_timestamp_ms = _now_in_milliseconds()
     if isinstance(path, str) or isinstance(path, epath.Path):
       path = os.fspath(path).split(",")
-    for p in path:
+    for p in path:  # pyrefly: ignore[not-iterable]
       for file in file_utils.expand_glob(p):
         self._info_proto.data_source_accesses.append(
             dataset_info_pb2.DataSourceAccess(
@@ -921,7 +921,7 @@ def __repr__(self):
       config_description = SKIP
 
     if self._info_proto.config_tags:
-      config_tags = ", ".join(self.config_tags)
+      config_tags = ", ".join(self.config_tags)  # pyrefly: ignore[no-matching-overload]
     else:
       config_tags = SKIP
 
@@ -1007,10 +1007,10 @@ def _nest_to_proto(nest: Nest) -> dataset_info_pb2.SupervisedKeys.Nest:
     for item in nest:
       proto.tuple.items.append(_nest_to_proto(item))
   elif nest_type is dict:
-    nest = {key: _nest_to_proto(value) for key, value in nest.items()}
-    proto.dict.CopyFrom(dataset_info_pb2.SupervisedKeys.Dict(dict=nest))
+    nest = {key: _nest_to_proto(value) for key, value in nest.items()}  # pyrefly: ignore[bad-assignment]
+    proto.dict.CopyFrom(dataset_info_pb2.SupervisedKeys.Dict(dict=nest))  # pyrefly: ignore[bad-argument-type]
   elif nest_type is str:
-    proto.feature_key = nest
+    proto.feature_key = nest  # pyrefly: ignore[bad-assignment]
   else:
     raise ValueError(
         "The nested structures in `supervised_keys` must only "
@@ -1072,7 +1072,7 @@ def _supervised_keys_from_proto(
   if proto.input and proto.output:
     return (proto.input, proto.output)
   elif proto.tuple:
-    return tuple(_nest_from_proto(item) for item in proto.tuple.items)
+    return tuple(_nest_from_proto(item) for item in proto.tuple.items)  # pyrefly: ignore[bad-return]
   else:
     raise ValueError(
         "A `SupervisedKeys` proto must have either `input` and "
@@ -1376,7 +1376,7 @@ def add_tfds_data_source_access(
               name=dataset_reference.dataset_name,
               config=dataset_reference.config,
               version=str(dataset_reference.version),
-              data_dir=os.fspath(dataset_reference.data_dir),
+              data_dir=os.fspath(dataset_reference.data_dir),  # pyrefly: ignore[no-matching-overload]
               ds_namespace=dataset_reference.namespace,
           ),
           url=dataset_info_pb2.Url(url=url),

tensorflow_datasets/core/dataset_metadata.py

@@ -88,8 +88,8 @@ def load(pkg_path: epath.Path) -> DatasetMetadata:
   raw_metadata = _read_files(pkg_path)
   tags = _get_tags(raw_metadata.get(TAGS_FILENAME, ""))
   return DatasetMetadata(
-      description=raw_metadata.get(DESCRIPTIONS_FILENAME, None),
-      citation=raw_metadata.get(CITATIONS_FILENAME, None),
+      description=raw_metadata.get(DESCRIPTIONS_FILENAME, None),  # pyrefly: ignore[bad-argument-type]
+      citation=raw_metadata.get(CITATIONS_FILENAME, None),  # pyrefly: ignore[bad-argument-type]
       tags=tags,
   )
 
@@ -108,4 +108,4 @@ def _read_files(path: epath.Path) -> dict[str, str]:
   for inode in path.iterdir():
     if inode.name in _METADATA_FILES:
       name2path[inode.name] = path.joinpath(inode.name)
-  return etree.parallel_map(lambda f: f.read_text(encoding="utf-8"), name2path)
+  return etree.parallel_map(lambda f: f.read_text(encoding="utf-8"), name2path)  # pyrefly: ignore[bad-return]

tensorflow_datasets/core/example_parser.py

@@ -130,8 +130,8 @@ def __post_init__(self):
   def parse_example(
       self, serialized_example: bytes | memoryview
   ) -> Mapping[str, Union[np.ndarray, list[Any]]]:
-    example = tf_example_pb2.Example.FromString(serialized_example)
-    np_example = _features_to_numpy(example.features, self._flat_example_specs)
+    example = tf_example_pb2.Example.FromString(serialized_example)  # pyrefly: ignore[bad-argument-type]
+    np_example = _features_to_numpy(example.features, self._flat_example_specs)  # pyrefly: ignore[bad-argument-type]
     return utils.pack_as_nest_dict(np_example, self.example_specs)
 
 
@@ -220,7 +220,7 @@ def _feature_to_numpy(
     return value_array.item()
   feature_name = _key_to_feature_name(key)
   row_lengths = _extract_row_lengths(features, feature_name)
-  return reshape_ragged_tensor(value_array, row_lengths, shape)
+  return reshape_ragged_tensor(value_array, row_lengths, shape)  # pyrefly: ignore[bad-argument-type]
 
 
 def _extract_row_lengths(
@@ -269,7 +269,7 @@ def reshape_ragged_tensor(
       value = values[i : i + length]
       array.append(value)
       i += length
-    values = array
+    values = array  # pyrefly: ignore[bad-assignment]
   return values
 
 
@@ -346,7 +346,7 @@ def _deserialize_single_field(
 
 def _dict_to_ragged(example_data, tensor_info):
   """Reconstruct the ragged tensor from the row ids."""
-  return tf.RaggedTensor.from_nested_row_lengths(
+  return tf.RaggedTensor.from_nested_row_lengths(  # pyrefly: ignore[missing-attribute]
       flat_values=example_data["ragged_flat_values"],
       nested_row_lengths=[
           example_data["ragged_row_lengths_{}".format(k)]
@@ -363,13 +363,13 @@ def _to_tf_example_spec(tensor_info: feature_lib.TensorInfo):
   # This create limitation like float64 downsampled to float32, bool converted
   # to int64 which is space ineficient, no support for complexes or quantized
   # It seems quite space inefficient to convert bool to int64
-  if dtype_utils.is_integer(tensor_info.tf_dtype) or dtype_utils.is_bool(
-      tensor_info.tf_dtype
+  if dtype_utils.is_integer(tensor_info.tf_dtype) or dtype_utils.is_bool(  # pyrefly: ignore[bad-argument-type]
+      tensor_info.tf_dtype  # pyrefly: ignore[bad-argument-type]
   ):
     dtype = tf.int64
-  elif dtype_utils.is_floating(tensor_info.tf_dtype):
+  elif dtype_utils.is_floating(tensor_info.tf_dtype):  # pyrefly: ignore[bad-argument-type]
     dtype = tf.float32
-  elif dtype_utils.is_string(tensor_info.tf_dtype):
+  elif dtype_utils.is_string(tensor_info.tf_dtype):  # pyrefly: ignore[bad-argument-type]
     dtype = tf.string
   else:
     # TFRecord only support 3 types

tensorflow_datasets/core/example_serializer.py

@@ -81,7 +81,7 @@ def get_tf_example(self, example: TreeDict[Any]) -> tf.train.Example:
     Returns:
       The `tf.train.Example` proto
     """
-    return _dict_to_tf_example(
+    return _dict_to_tf_example(  # pyrefly: ignore[bad-return]
         utils.flatten_nest_dict(example), self._flat_example_specs
     )
 
@@ -166,7 +166,7 @@ def _is_string(item) -> bool:
 def _item_to_np_array(item, dtype: np.dtype, shape: Shape) -> np.ndarray:
   """Single item to a np.array."""
   result = np.asanyarray(item, dtype=dtype)
-  utils.assert_shape_match(result.shape, shape)
+  utils.assert_shape_match(result.shape, shape)  # pyrefly: ignore[bad-argument-type]
   if dtype_utils.is_string(dtype) and not _is_string(item):
     raise ValueError(
         f"Unsupported value: {result}\nCould not convert to bytes list."
@@ -192,16 +192,16 @@ def _item_to_tf_feature(
     v = v.view(np.int64)
   vals = v.flat  # Convert v into a 1-d array (without extra copy)
   if dtype_utils.is_integer(v.dtype):
-    return tf_feature_pb2.Feature(
-        int64_list=tf_feature_pb2.Int64List(value=vals)
+    return tf_feature_pb2.Feature(  # pyrefly: ignore[bad-return]
+        int64_list=tf_feature_pb2.Int64List(value=vals)  # pyrefly: ignore[bad-argument-type]
     )
   elif dtype_utils.is_floating(v.dtype):
-    return tf_feature_pb2.Feature(
-        float_list=tf_feature_pb2.FloatList(value=vals)
+    return tf_feature_pb2.Feature(  # pyrefly: ignore[bad-return]
+        float_list=tf_feature_pb2.FloatList(value=vals)  # pyrefly: ignore[bad-argument-type]
     )
   elif dtype_utils.is_string(tensor_info.np_dtype):
-    vals = [_as_bytes(x) for x in vals]
-    return tf_feature_pb2.Feature(
+    vals = [_as_bytes(x) for x in vals]  # pyrefly: ignore[bad-argument-type]
+    return tf_feature_pb2.Feature(  # pyrefly: ignore[bad-return]
         bytes_list=tf_feature_pb2.BytesList(value=vals)
     )
   else:
@@ -229,7 +229,7 @@ def _as_bytes(bytes_or_text: bytearray | bytes | str) -> bytes:
   """
   # Validate encoding, a LookupError will be raised if invalid.
   if isinstance(bytes_or_text, (bytearray, bytes)):
-    return bytes_or_text
+    return bytes_or_text  # pyrefly: ignore[bad-return]
   elif isinstance(bytes_or_text, str):
     return bytes_or_text.encode(encoding="utf-8")
   else:
@@ -300,11 +300,11 @@ def _add_ragged_fields(example_data, tensor_info: feature_lib.TensorInfo):
     tensor_info_length = feature_lib.TensorInfo(shape=(None,), dtype=np.int64)
     ragged_attr_dict = {
         "ragged_row_lengths_{}".format(i): (length, tensor_info_length)
-        for i, length in enumerate(nested_row_lengths)
+        for i, length in enumerate(nested_row_lengths)  # pyrefly: ignore[unbound-name]
     }
     tensor_info_flat = feature_lib.TensorInfo(
-        shape=(None,) + tensor_info.shape[tensor_info.sequence_rank :],
-        dtype=tensor_info.dtype,
+        shape=(None,) + tensor_info.shape[tensor_info.sequence_rank :],  # pyrefly: ignore[unsupported-operation]
+        dtype=tensor_info.dtype,  # pyrefly: ignore[bad-argument-type]
     )
     ragged_attr_dict["ragged_flat_values"] = (example_data, tensor_info_flat)
     return ragged_attr_dict
@@ -342,8 +342,8 @@ def _extract_ragged_attributes(
       )
   )
   if not flat_values:  # The full sequence is empty
-    flat_values = np.empty(
-        shape=(0,) + tensor_info.shape[tensor_info.sequence_rank :],
+    flat_values = np.empty(  # pyrefly: ignore[no-matching-overload]
+        shape=(0,) + tensor_info.shape[tensor_info.sequence_rank :],  # pyrefly: ignore[unsupported-operation]
         dtype=tensor_info.np_dtype,
     )
   else:  # Otherwise, merge all flat values together, some might be empty

tensorflow_datasets/core/features/audio_feature.py

@@ -179,11 +179,11 @@ def decode_audio(self, audio_tensor: tf.Tensor) -> tf.Tensor:
       decoded_audio_tensor = tf.squeeze(decoded_audio_tensor)
     else:
       if self._file_format:
-        file_format_tensor = tf.experimental.Optional.from_value(
+        file_format_tensor = tf.experimental.Optional.from_value(  # pyrefly: ignore[missing-attribute]
             self._file_format
         )
       else:
-        file_format_tensor = tf.experimental.Optional.empty(
+        file_format_tensor = tf.experimental.Optional.empty(  # pyrefly: ignore[missing-attribute]
             tf.TensorSpec(shape=(), dtype=tf.string)
         )
 
@@ -206,7 +206,7 @@ class _EagerDecoder(_AudioDecoder):
   def encode_audio(
       self, fobj: BinaryIO, file_format: Optional[str]
   ) -> np.ndarray:
-    audio = _pydub_load_audio(fobj, file_format, self._channels)
+    audio = _pydub_load_audio(fobj, file_format, self._channels)  # pyrefly: ignore[bad-argument-type]
     return audio.astype(self._np_dtype)
 
   def decode_audio(self, audio_tensor: tf.Tensor) -> tf.Tensor:
@@ -336,17 +336,17 @@ def repr_html(self, ex: np.ndarray) -> str:
     )
 
   @classmethod
-  def from_json_content(
+  def from_json_content(  # pyrefly: ignore[bad-override]
       cls, value: Union[Json, feature_pb2.AudioFeature]
   ) -> 'Audio':
     if isinstance(value, dict):
       # For backwards compatibility
       return cls(
-          file_format=value['file_format'],
-          shape=tuple(value['shape']),
-          dtype=feature_lib.dtype_from_str(value['dtype']),
-          sample_rate=value['sample_rate'],
-          lazy_decode=value.get('lazy_decode', False),
+          file_format=value['file_format'],  # pyrefly: ignore[bad-argument-type]
+          shape=tuple(value['shape']),  # pyrefly: ignore[bad-argument-type]
+          dtype=feature_lib.dtype_from_str(value['dtype']),  # pyrefly: ignore[bad-argument-type]
+          sample_rate=value['sample_rate'],  # pyrefly: ignore[bad-argument-type]
+          lazy_decode=value.get('lazy_decode', False),  # pyrefly: ignore[bad-argument-type]
       )
     return cls(
         shape=feature_lib.from_shape_proto(value.shape),
@@ -360,7 +360,7 @@ def from_json_content(
   def to_json_content(self) -> feature_pb2.AudioFeature:  # pytype: disable=signature-mismatch  # overriding-return-type-checks
     return feature_pb2.AudioFeature(
         shape=feature_lib.to_shape_proto(self.shape),
-        dtype=feature_lib.dtype_to_str(self.dtype),
+        dtype=feature_lib.dtype_to_str(self.dtype),  # pyrefly: ignore[bad-argument-type]
         file_format=self._file_format,
         sample_rate=self._sample_rate,
         encoding=self._encoding.name,