From fa62a1b4491c86ffb285d525e4049446bc30563a Mon Sep 17 00:00:00 2001 From: Hana Joo Date: Tue, 7 Jul 2026 17:39:34 -0700 Subject: [PATCH] Automated Code Change PiperOrigin-RevId: 944182406 --- tensorflow_datasets/core/as_dataframe.py | 4 +- tensorflow_datasets/core/dataset_builder.py | 20 +++++----- .../core/dataset_collection_builder.py | 2 +- tensorflow_datasets/core/dataset_info.py | 20 +++++----- tensorflow_datasets/core/dataset_metadata.py | 6 +-- tensorflow_datasets/core/example_parser.py | 18 ++++----- .../core/example_serializer.py | 28 +++++++------- .../core/features/audio_feature.py | 20 +++++----- .../core/features/bounding_boxes.py | 6 +-- .../core/features/bounding_boxes_utils.py | 2 +- .../core/features/class_label_feature.py | 16 ++++---- tensorflow_datasets/core/features/feature.py | 24 ++++++------ .../core/features/features_dict.py | 10 ++--- .../core/features/image_feature.py | 38 +++++++++---------- .../core/features/labeled_image.py | 4 +- tensorflow_datasets/core/features/scalar.py | 2 +- .../core/features/sequence_feature.py | 8 ++-- .../core/features/tensor_feature.py | 28 +++++++------- .../core/features/text_feature.py | 18 ++++----- .../core/features/top_level_feature.py | 2 +- .../core/features/translation_feature.py | 12 +++--- .../core/features/video_feature.py | 16 ++++---- tensorflow_datasets/core/file_adapters.py | 2 +- tensorflow_datasets/core/load.py | 14 +++---- tensorflow_datasets/core/naming.py | 6 +-- tensorflow_datasets/core/read_only_builder.py | 23 +++++------ tensorflow_datasets/core/reader.py | 20 +++++----- tensorflow_datasets/core/sequential_writer.py | 2 +- tensorflow_datasets/core/shuffle.py | 8 ++-- tensorflow_datasets/core/split_builder.py | 16 ++++---- tensorflow_datasets/core/splits.py | 20 +++++----- tensorflow_datasets/core/subsplits_utils.py | 2 +- tensorflow_datasets/core/writer.py | 6 +-- 33 files changed, 212 insertions(+), 211 deletions(-) diff --git a/tensorflow_datasets/core/as_dataframe.py b/tensorflow_datasets/core/as_dataframe.py index 08f01da6b9e..4536e961bfd 100644 --- a/tensorflow_datasets/core/as_dataframe.py +++ b/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}) diff --git a/tensorflow_datasets/core/dataset_builder.py b/tensorflow_datasets/core/dataset_builder.py index a1ad6a1fe44..02a21b3d695 100644 --- a/tensorflow_datasets/core/dataset_builder.py +++ b/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: diff --git a/tensorflow_datasets/core/dataset_collection_builder.py b/tensorflow_datasets/core/dataset_collection_builder.py index dec72ce5385..4a0d62faa47 100644 --- a/tensorflow_datasets/core/dataset_collection_builder.py +++ b/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, ) diff --git a/tensorflow_datasets/core/dataset_info.py b/tensorflow_datasets/core/dataset_info.py index 1bb51b01838..38fcf8404eb 100644 --- a/tensorflow_datasets/core/dataset_info.py +++ b/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), diff --git a/tensorflow_datasets/core/dataset_metadata.py b/tensorflow_datasets/core/dataset_metadata.py index 68126f322bc..6943769c3e1 100644 --- a/tensorflow_datasets/core/dataset_metadata.py +++ b/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] diff --git a/tensorflow_datasets/core/example_parser.py b/tensorflow_datasets/core/example_parser.py index 7f0c6a33ebb..92a295bf801 100644 --- a/tensorflow_datasets/core/example_parser.py +++ b/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 diff --git a/tensorflow_datasets/core/example_serializer.py b/tensorflow_datasets/core/example_serializer.py index d7f426de24d..51214732239 100644 --- a/tensorflow_datasets/core/example_serializer.py +++ b/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 diff --git a/tensorflow_datasets/core/features/audio_feature.py b/tensorflow_datasets/core/features/audio_feature.py index ff364c444c8..ec01ffea117 100644 --- a/tensorflow_datasets/core/features/audio_feature.py +++ b/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, diff --git a/tensorflow_datasets/core/features/bounding_boxes.py b/tensorflow_datasets/core/features/bounding_boxes.py index de14933c3ec..e74a915e875 100644 --- a/tensorflow_datasets/core/features/bounding_boxes.py +++ b/tensorflow_datasets/core/features/bounding_boxes.py @@ -139,18 +139,18 @@ def repr_html_batch(self, ex: np.ndarray) -> str: return _repr_html(ex, bbox_format=self.bbox_format) @classmethod - def from_json_content( + def from_json_content( # pyrefly: ignore[bad-override] cls, value: Union[Json, feature_pb2.BoundingBoxFeature] ) -> 'BBoxFeature': if isinstance(value, dict): - return cls(**value) + return cls(**value) # pyrefly: ignore[bad-argument-type] return cls( bbox_format=bb_utils.BBoxFormat(value.bbox_format) if value.bbox_format else None ) - def to_json_content( + def to_json_content( # pyrefly: ignore[bad-override] self, ) -> ( feature_pb2.BoundingBoxFeature diff --git a/tensorflow_datasets/core/features/bounding_boxes_utils.py b/tensorflow_datasets/core/features/bounding_boxes_utils.py index 5bac649a5b4..bc4bdb001f2 100644 --- a/tensorflow_datasets/core/features/bounding_boxes_utils.py +++ b/tensorflow_datasets/core/features/bounding_boxes_utils.py @@ -123,7 +123,7 @@ def convert_coordinates_to_bbox( ) if normalize: - bbox = normalize_bbox(bbox, img_shape) + bbox = normalize_bbox(bbox, img_shape) # pyrefly: ignore[bad-argument-type] return bbox diff --git a/tensorflow_datasets/core/features/class_label_feature.py b/tensorflow_datasets/core/features/class_label_feature.py index 67e7128ab4d..fedf259e3fc 100644 --- a/tensorflow_datasets/core/features/class_label_feature.py +++ b/tensorflow_datasets/core/features/class_label_feature.py @@ -88,7 +88,7 @@ def num_classes(self) -> Optional[int]: @property def names(self) -> list[str]: if not self._int2str: - return [str(i) for i in range(self._num_classes)] + return [str(i) for i in range(self._num_classes)] # pyrefly: ignore[bad-argument-type] return self._int2str @names.setter @@ -132,7 +132,7 @@ def str2int(self, str_value: str) -> int: int_value = int(str_value) except ValueError: failed_parse = True - if failed_parse or not 0 <= int_value < self._num_classes: + if failed_parse or not 0 <= int_value < self._num_classes: # pyrefly: ignore[unbound-name, unsupported-operation] raise ValueError("Invalid string class label %s" % str_value) return int_value @@ -146,7 +146,7 @@ def int2str(self, int_value: int) -> str: return self._int2str[int_value] # No names provided, return str(int) - if not 0 <= int_value < self._num_classes: + if not 0 <= int_value < self._num_classes: # pyrefly: ignore[unsupported-operation] raise ValueError("Invalid integer class label %d" % int_value) return str(int_value) @@ -177,13 +177,13 @@ def save_metadata(self, data_dir, feature_name=None) -> None: """See base class for details.""" # Save names if defined if self._str2int is not None: - names_filepath = self.get_names_filepath(data_dir, feature_name) + names_filepath = self.get_names_filepath(data_dir, feature_name) # pyrefly: ignore[bad-argument-type] _write_names_to_file(names_filepath, self.names) def load_metadata(self, data_dir, feature_name=None) -> Optional[list[str]]: """See base class for details.""" # Restore names if defined - names_filepath = self.get_names_filepath(data_dir, feature_name) + names_filepath = self.get_names_filepath(data_dir, feature_name) # pyrefly: ignore[bad-argument-type] try: self.names = _load_names_from_file(names_filepath) except OSError: @@ -192,7 +192,7 @@ def load_metadata(self, data_dir, feature_name=None) -> Optional[list[str]]: def _additional_repr_info(self) -> dict[str, int]: return {"num_classes": self.num_classes} # pytype: disable=bad-return-type # always-use-property-annotation - def repr_html(self, ex: int) -> str: + def repr_html(self, ex: int) -> str: # pyrefly: ignore[bad-override] """Class labels are displayed with their name.""" if ex == -1: return "-" @@ -202,11 +202,11 @@ def repr_html(self, ex: int) -> str: return f"{ex} ({self.int2str(ex)})" @classmethod - def from_json_content( + def from_json_content( # pyrefly: ignore[bad-override] cls, value: Union[Json, feature_pb2.ClassLabel] ) -> "ClassLabel": if isinstance(value, dict): - return cls(**value) + return cls(**value) # pyrefly: ignore[bad-argument-type] return cls(num_classes=value.num_classes) def to_json_content(self) -> feature_pb2.ClassLabel: # pytype: disable=signature-mismatch # overriding-return-type-checks diff --git a/tensorflow_datasets/core/features/feature.py b/tensorflow_datasets/core/features/feature.py index 4dcfd036176..a8ee6d0c0c4 100644 --- a/tensorflow_datasets/core/features/feature.py +++ b/tensorflow_datasets/core/features/feature.py @@ -180,8 +180,8 @@ def to_tensor_spec(self) -> tf.TensorSpec: dtype = self.tf_dtype shape = _to_tensor_shape(self.shape) if self.dataset_lvl > 1 or self.sequence_rank > 1: - return tf.RaggedTensorSpec(dtype=dtype, shape=shape) - return tf.TensorSpec(dtype=dtype, shape=shape) + return tf.RaggedTensorSpec(dtype=dtype, shape=shape) # pyrefly: ignore[bad-argument-type, bad-return] + return tf.TensorSpec(dtype=dtype, shape=shape) # pyrefly: ignore[bad-argument-type] def to_bounded_tensor_spec(self) -> tf_agents.specs.BoundedTensorSpec: """Converts this TensorInfo instance to a tf.BoundedTensorSpec. @@ -456,10 +456,10 @@ def from_json(cls, value: Json) -> FeatureConnector: if 'type' in value: # Legacy mode class_name = value['type'] # my_project.xyz.MyFeature content = value['content'] - feature_cls = cls.cls_from_name(class_name) + feature_cls = cls.cls_from_name(class_name) # pyrefly: ignore[bad-argument-type] proto_cls_name = value.get('proto_cls') if proto_cls_name: # The content is a proto, need to reconstruct it - proto_cls = _name2proto_cls(proto_cls_name) + proto_cls = _name2proto_cls(proto_cls_name) # pyrefly: ignore[bad-argument-type] if isinstance(content, str): # Backward compatible mode content = json_format.Parse(content, proto_cls()) elif isinstance(content, dict): @@ -472,8 +472,8 @@ def from_json(cls, value: Json) -> FeatureConnector: if isinstance(content, dict) and ( 'dtype' in content and isinstance(content['dtype'], str) ): - content['dtype'] = dtype_from_str(content['dtype']) - return feature_cls.from_json_content(content) + content['dtype'] = dtype_from_str(content['dtype']) # pyrefly: ignore[unsupported-operation] + return feature_cls.from_json_content(content) # pyrefly: ignore[bad-argument-type] else: feature_proto = json_format.ParseDict(value, feature_pb2.Feature()) return FeatureConnector.from_proto(feature_proto) @@ -611,7 +611,7 @@ def to_proto(self) -> feature_pb2.Feature: python_class_name=self._fully_qualified_class_name, description=self._doc.desc, value_range=self._doc.value_range, - **oneof_kwarg, + **oneof_kwarg, # pyrefly: ignore[bad-argument-type] ) @classmethod @@ -628,7 +628,7 @@ def from_proto(cls, feature_proto: feature_pb2.Feature) -> T: # Not all feature classes accept the documentation as an argument. feature = feature_cls.from_json_content(value=feature_content) feature._set_doc(Documentation.from_proto(feature_proto)) # pylint: disable=protected-access - return feature + return feature # pyrefly: ignore[bad-return] def save_config(self, root_dir: str) -> None: """Exports the `FeatureConnector` to a file. @@ -695,7 +695,7 @@ def get_serialized_info(self) -> Union[TensorInfo, Mapping[str, TensorInfo]]: Returns: features: Either a dict of feature proto object, or a feature proto object """ - return self.get_tensor_info() + return self.get_tensor_info() # pyrefly: ignore[bad-return] @abc.abstractmethod def encode_example(self, example_data): @@ -821,7 +821,7 @@ def decode_batch_example(self, tfexample_data): # `(0,)` -> `(0, None, None, 3)`. # Instead, we arbitrarily set unknown shape to `0`: # `(0,)` -> `(0, 0, 0, 3)` - tf_type = tf.dtypes.as_dtype(self.dtype) + tf_type = tf.dtypes.as_dtype(self.dtype) # pyrefly: ignore[bad-argument-type] return tf.cond( tf.equal(tf.shape(ex)[0], 0), # Empty sequence lambda: _make_empty_seq_output(shape=self.shape, dtype=tf_type), @@ -1023,8 +1023,8 @@ def catalog_documentation(self) -> List[CatalogFeatureDocumentation]: name=feature_name, cls_name=type(self).__name__, tensor_info=tensor_info, - description=self._doc.desc, - value_range=self._doc.value_range, + description=self._doc.desc, # pyrefly: ignore[bad-argument-type] + value_range=self._doc.value_range, # pyrefly: ignore[bad-argument-type] ) ) return result diff --git a/tensorflow_datasets/core/features/features_dict.py b/tensorflow_datasets/core/features/features_dict.py index b77ac59ba28..944a918a24a 100644 --- a/tensorflow_datasets/core/features/features_dict.py +++ b/tensorflow_datasets/core/features/features_dict.py @@ -179,8 +179,8 @@ def catalog_documentation( name='', cls_name=type(self).__name__, tensor_info=None, - description=self._doc.desc, - value_range=self._doc.value_range, + description=self._doc.desc, # pyrefly: ignore[bad-argument-type] + value_range=self._doc.value_range, # pyrefly: ignore[bad-argument-type] ) ] for feature_name, feature in sorted(list(self._feature_dict.items())): @@ -209,12 +209,12 @@ def get_serialized_info(self): } @classmethod - def from_json_content( + def from_json_content( # pyrefly: ignore[bad-override] cls, value: Json | feature_pb2.FeaturesDict ) -> 'FeaturesDict': if isinstance(value, dict): features = { - k: feature_lib.FeatureConnector.from_json(v) for k, v in value.items() + k: feature_lib.FeatureConnector.from_json(v) for k, v in value.items() # pyrefly: ignore[bad-argument-type] } else: features = { @@ -315,7 +315,7 @@ def to_feature(value: feature_lib.FeatureConnectorArg): if isinstance(value, feature_lib.FeatureConnector): return value elif dtype_utils.is_np_or_tf_dtype(value): # tf.int32, np.int32,... - return tensor_feature.Tensor(shape=(), dtype=value) + return tensor_feature.Tensor(shape=(), dtype=value) # pyrefly: ignore[bad-argument-type] elif isinstance(value, dict): return FeaturesDict(value) else: diff --git a/tensorflow_datasets/core/features/image_feature.py b/tensorflow_datasets/core/features/image_feature.py index 5c2df84e499..c19f7766e5c 100644 --- a/tensorflow_datasets/core/features/image_feature.py +++ b/tensorflow_datasets/core/features/image_feature.py @@ -124,14 +124,14 @@ def encode_image_or_path(self, image_or_path_or_fobj): def _encode_image(self, np_image: np.ndarray) -> bytes: """Returns np_image encoded as jpeg or png.""" - _validate_np_array(np_image, shape=self.shape, dtype=self.np_dtype) + _validate_np_array(np_image, shape=self.shape, dtype=self.np_dtype) # pyrefly: ignore[bad-argument-type] # Should we be more strict about explicitly define the encoding (raise # error / warning instead) ? # It has created subtle issues for imagenet_corrupted: images are read as # JPEG images to apply some processing, but final image saved as PNG # (default) rather than JPEG. - return self._runner.run(_ENCODE_FN[self.encoding_format](), np_image) + return self._runner.run(_ENCODE_FN[self.encoding_format](), np_image) # pyrefly: ignore[bad-index] def _encode_pil_image(self, pil_image) -> bytes: """Encode a PIL Image object to bytes. @@ -149,7 +149,7 @@ def _encode_pil_image(self, pil_image) -> bytes: def decode_image(self, img: tf.Tensor) -> tf.Tensor: """Decode the jpeg or png bytes to 3d tensor.""" - tf_dtype = tf.dtypes.as_dtype(self.dtype) + tf_dtype = tf.dtypes.as_dtype(self.dtype) # pyrefly: ignore[bad-argument-type] img = tf.image.decode_image(img, channels=self.shape[-1], dtype=tf_dtype) img.set_shape(self.shape) return img @@ -180,7 +180,7 @@ def __init__( ) self._float_shape = shape super().__init__( - shape=shape[:2] + (4,), + shape=shape[:2] + (4,), # pyrefly: ignore[unsupported-operation] dtype=np.uint8, encoding_format=encoding_format, ) @@ -198,8 +198,8 @@ def _encode_image(self, np_image: np.ndarray) -> bytes: _validate_np_array(np_image, shape=self._float_shape, dtype=np.float32) # Bitcast 1 channel float32 -> 4 channels uint8 np_image = np_image.view(np.uint8) - np_image = super()._encode_image(np_image) - return np_image + np_image = super()._encode_image(np_image) # pyrefly: ignore[bad-assignment] + return np_image # pyrefly: ignore[bad-return] def decode_image(self, img: tf.Tensor) -> tf.Tensor: img = super().decode_image(img) @@ -295,7 +295,7 @@ def __init__( if self._dtype == np.float32: # Float images encoded as 4-channels uint8 self._image_encoder = _FloatImageEncoder( shape=self._shape, - encoding_format=self._encoding_format, + encoding_format=self._encoding_format, # pyrefly: ignore[bad-argument-type] ) else: self._image_encoder = _ImageEncoder( @@ -359,16 +359,16 @@ def decode_example_np_with_opencv( ) -> np.ndarray: """Reconstruct the image with OpenCV from bytes.""" assert cv2, 'OpenCV is not installed. OpenCV is required for this method.' - example = np.frombuffer(example, dtype=np.uint8) + example = np.frombuffer(example, dtype=np.uint8) # pyrefly: ignore[bad-assignment] example = cv2.imdecode(example, cv2.IMREAD_UNCHANGED) dtype = self.np_dtype if self.np_dtype != np.float32 else np.uint8 example = example.astype(dtype, copy=False) - example = _reorder_opencv_channels(example) - example = _reshape_grayscale_image(example, num_channels) + example = _reorder_opencv_channels(example) # pyrefly: ignore[bad-assignment] + example = _reshape_grayscale_image(example, num_channels) # pyrefly: ignore[bad-argument-type, bad-assignment] # Bitcast 4 channels uint8 -> 1 channel float32. if self.np_dtype == np.float32: - return example.view(np.float32) - return example + return example.view(np.float32) # pyrefly: ignore[missing-attribute] + return example # pyrefly: ignore[bad-return] def decode_example_np_with_pil( self, example: bytes, num_channels: int @@ -380,7 +380,7 @@ def decode_example_np_with_pil( bytes_io = io.BytesIO(example) with PIL_Image.open(bytes_io) as image: dtype = self.np_dtype if self.np_dtype != np.float32 else np.uint8 - image = np.asarray(image, dtype=dtype) + image = np.asarray(image, dtype=dtype) # pyrefly: ignore[no-matching-overload] image = _reshape_grayscale_image(image, num_channels) if self.np_dtype == np.uint8: return image @@ -412,16 +412,16 @@ def repr_html_batch(self, ex: np.ndarray) -> str: return make_video_repr_html(ex, use_colormap=self._use_colormap) @classmethod - def from_json_content( + def from_json_content( # pyrefly: ignore[bad-override] cls, value: Union[Json, feature_pb2.ImageFeature] ) -> 'Image': if isinstance(value, dict): # For backwards compatibility return cls( # pytype: disable=wrong-arg-types - shape=tuple(value['shape']), - dtype=feature_lib.dtype_from_str(value['dtype']), - encoding_format=value['encoding_format'], - use_colormap=value.get('use_colormap'), + shape=tuple(value['shape']), # pyrefly: ignore[bad-argument-type] + dtype=feature_lib.dtype_from_str(value['dtype']), # pyrefly: ignore[bad-argument-type] + encoding_format=value['encoding_format'], # pyrefly: ignore[bad-argument-type] + use_colormap=value.get('use_colormap'), # pyrefly: ignore[bad-argument-type] ) return cls( shape=feature_lib.from_shape_proto(value.shape), @@ -594,7 +594,7 @@ def _validate_np_array( raise ValueError( f'Image dtype should be {dtype}. Detected: {np_array.dtype}.' ) - utils.assert_shape_match(np_array.shape, shape) + utils.assert_shape_match(np_array.shape, shape) # pyrefly: ignore[bad-argument-type] def _reorder_opencv_channels(example: np.ndarray) -> np.ndarray: diff --git a/tensorflow_datasets/core/features/labeled_image.py b/tensorflow_datasets/core/features/labeled_image.py index 587c88b1fa0..695089b53e4 100644 --- a/tensorflow_datasets/core/features/labeled_image.py +++ b/tensorflow_datasets/core/features/labeled_image.py @@ -87,7 +87,7 @@ def __init__( f'LabeledImage shape should have a single channel. Got: {shape}' ) label_kwargs = _labels_to_kwarg(labels) - self._class_label = class_label_feature.ClassLabel(**label_kwargs) + self._class_label = class_label_feature.ClassLabel(**label_kwargs) # pyrefly: ignore[bad-argument-type] @property def num_classes(self) -> Optional[int]: @@ -118,7 +118,7 @@ def from_json_content( ) -> 'LabeledImage': if isinstance(value, dict): # For backwards compatibility - return cls(**value) + return cls(**value) # pyrefly: ignore[bad-argument-type] return cls( shape=feature_lib.from_shape_proto(value.shape), dtype=feature_lib.dtype_from_str(value.dtype), diff --git a/tensorflow_datasets/core/features/scalar.py b/tensorflow_datasets/core/features/scalar.py index 12b5efcf896..22b366352ef 100644 --- a/tensorflow_datasets/core/features/scalar.py +++ b/tensorflow_datasets/core/features/scalar.py @@ -39,7 +39,7 @@ def __init__( super().__init__(shape=(), dtype=dtype, doc=doc, optional=optional) @classmethod - def from_json_content(cls, value: feature_pb2.TensorFeature) -> 'Scalar': + def from_json_content(cls, value: feature_pb2.TensorFeature) -> 'Scalar': # pyrefly: ignore[bad-override] return cls( dtype=feature_lib.dtype_from_str(value.dtype), optional=value.optional, diff --git a/tensorflow_datasets/core/features/sequence_feature.py b/tensorflow_datasets/core/features/sequence_feature.py index 1448557a794..760ac9cbbc9 100644 --- a/tensorflow_datasets/core/features/sequence_feature.py +++ b/tensorflow_datasets/core/features/sequence_feature.py @@ -112,7 +112,7 @@ def feature(self): def _add_length_dim(self, tensor_info): """Add the length dimension to the given tensor_info.""" tensor_info = feature_lib.TensorInfo.copy_from(tensor_info) - tensor_info.shape = (self._length,) + tensor_info.shape + tensor_info.shape = (self._length,) + tensor_info.shape # pyrefly: ignore[unsupported-operation] tensor_info.sequence_rank += 1 return tensor_info @@ -234,15 +234,15 @@ def catalog_documentation( return result @classmethod - def from_json_content( + def from_json_content( # pyrefly: ignore[bad-override] cls, value: Union[Json, feature_pb2.Sequence], ) -> 'Sequence': if isinstance(value, dict): # For backwards compatibility return cls( - feature=feature_lib.FeatureConnector.from_json(value['feature']), - length=value['length'], + feature=feature_lib.FeatureConnector.from_json(value['feature']), # pyrefly: ignore[bad-argument-type] + length=value['length'], # pyrefly: ignore[bad-argument-type] ) return cls( feature=feature_lib.FeatureConnector.from_proto(value.feature), diff --git a/tensorflow_datasets/core/features/tensor_feature.py b/tensorflow_datasets/core/features/tensor_feature.py index 04985a8c169..cc5f32c5094 100644 --- a/tensorflow_datasets/core/features/tensor_feature.py +++ b/tensorflow_datasets/core/features/tensor_feature.py @@ -189,16 +189,16 @@ def _default_value(self) -> Any: if dtype_utils.is_string(dtype): default_value = b'' elif dtype_utils.is_integer(dtype): - default_value = np.iinfo(dtype).min + default_value = np.iinfo(dtype).min # pyrefly: ignore[no-matching-overload] elif dtype_utils.is_floating(dtype): - default_value = np.finfo(dtype).min + default_value = np.finfo(dtype).min # pyrefly: ignore[no-matching-overload] elif dtype_utils.is_bool(dtype): default_value = False else: raise ValueError(f'Unsupported dtype: {dtype}') if not self._shape: return default_value - return np.full(self._shape, default_value, dtype=dtype) + return np.full(self._shape, default_value, dtype=dtype) # pyrefly: ignore[no-matching-overload] def encode_example(self, example_data): """See base class for details.""" @@ -346,18 +346,18 @@ def decode_ragged_example_np(self, example_data): return self.decode_example_np(example_data) @classmethod - def from_json_content( + def from_json_content( # pyrefly: ignore[bad-override] cls, value: Union[Json, feature_pb2.TensorFeature] ) -> 'Tensor': if isinstance(value, dict): return cls( - shape=tuple(value['shape']), - dtype=feature_lib.dtype_from_str(value['dtype']), + shape=tuple(value['shape']), # pyrefly: ignore[bad-argument-type] + dtype=feature_lib.dtype_from_str(value['dtype']), # pyrefly: ignore[bad-argument-type] # Use .get for backward-compatibility - encoding=value.get('encoding', Encoding.NONE), - minimum=value.get('minimum', None), - maximum=value.get('maximum', None), - optional=value.get('optional', False), + encoding=value.get('encoding', Encoding.NONE), # pyrefly: ignore[bad-argument-type] + minimum=value.get('minimum', None), # pyrefly: ignore[bad-argument-type] + maximum=value.get('maximum', None), # pyrefly: ignore[bad-argument-type] + optional=value.get('optional', False), # pyrefly: ignore[bad-argument-type] ) return cls( shape=feature_lib.from_shape_proto(value.shape), @@ -373,8 +373,8 @@ def to_json_content(self) -> feature_pb2.TensorFeature: shape=feature_lib.to_shape_proto(self._shape), dtype=feature_lib.dtype_to_str(self._dtype), encoding=self._encoding.value, - minimum=self._minimum, - maximum=self._maximum, + minimum=self._minimum, # pyrefly: ignore[bad-argument-type] + maximum=self._maximum, # pyrefly: ignore[bad-argument-type] optional=self._optional, ) @@ -403,7 +403,7 @@ def _execute_function_on_array_or_scalar( return function(data, *args, **kwargs) if isinstance(data, np.ndarray): partial_function = functools.partial(function, *args, **kwargs) - return np.array(list(map(partial_function, data))) + return np.array(list(map(partial_function, data))) # pyrefly: ignore[bad-return] raise ValueError( 'example should have type `bytes` or `np.ndarray(dtype=bytes)`, but' f' has wrong type {type(data)}' @@ -427,6 +427,6 @@ def get_inner_feature_repr(feature): # * For the base `Tensor` class (and not subclass). # * When shape is scalar (explicit check to avoid trigger when `shape=None`). if type(feature) == Tensor and feature.shape == (): # pylint: disable=unidiomatic-typecheck,g-explicit-bool-comparison - return feature_lib.dtype_to_str(feature.np_dtype) + return feature_lib.dtype_to_str(feature.np_dtype) # pyrefly: ignore[bad-argument-type] else: return repr(feature) diff --git a/tensorflow_datasets/core/features/text_feature.py b/tensorflow_datasets/core/features/text_feature.py index b5bab2126c0..d2a3861cdc0 100644 --- a/tensorflow_datasets/core/features/text_feature.py +++ b/tensorflow_datasets/core/features/text_feature.py @@ -130,13 +130,13 @@ def encode_example(self, example_data): def decode_example_np(self, example_data): return example_data - def save_metadata(self, data_dir, feature_name: str) -> None: + def save_metadata(self, data_dir, feature_name: str) -> None: # pyrefly: ignore[bad-override] if not self.encoder: return fname_prefix = _file_name_prefix_for_metadata(feature_name, data_dir) self.encoder.save_to_file(fname_prefix) - def load_metadata(self, data_dir, feature_name: str) -> None: + def load_metadata(self, data_dir, feature_name: str) -> None: # pyrefly: ignore[bad-override] if self._encoder_cls: fname_prefix = _file_name_prefix_for_metadata(feature_name, data_dir) self._encoder = self._encoder_cls.load_from_file(fname_prefix) # pytype: disable=attribute-error @@ -168,7 +168,7 @@ def maybe_build_from_corpus(self, corpus_generator, **kwargs): if self.encoder: return - vocab_size = self._encoder_config.vocab_size + vocab_size = self._encoder_config.vocab_size # pyrefly: ignore[missing-attribute] self.encoder = text_lib.SubwordTextEncoder.build_from_corpus( corpus_generator=corpus_generator, target_vocab_size=vocab_size, @@ -184,22 +184,22 @@ def _additional_repr_info(self): return {} return {"encoder": repr(self.encoder)} - def repr_html(self, ex: bytes) -> str: + def repr_html(self, ex: bytes) -> str: # pyrefly: ignore[bad-override] """Text are decoded.""" if self.encoder is not None: return repr(ex) try: - ex = ex.decode("utf-8") + ex = ex.decode("utf-8") # pyrefly: ignore[bad-assignment] except UnicodeDecodeError: # Some datasets have invalid UTF-8 examples (e.g. opinosis) return repr(ex[:1000]) - ex = html.escape(ex) - ex = textwrap.shorten(ex, width=1000) # Truncate long text - return ex + ex = html.escape(ex) # pyrefly: ignore[bad-argument-type, bad-assignment] + ex = textwrap.shorten(ex, width=1000) # Truncate long text # pyrefly: ignore[bad-argument-type, bad-assignment] + return ex # pyrefly: ignore[bad-return] @classmethod - def from_json_content( + def from_json_content( # pyrefly: ignore[bad-override] cls, value: Union[Json, feature_pb2.TextFeature] ) -> "Text": if isinstance(value, dict) and "use_encoder" in value: diff --git a/tensorflow_datasets/core/features/top_level_feature.py b/tensorflow_datasets/core/features/top_level_feature.py index 706732833cb..9d01fa35a93 100644 --- a/tensorflow_datasets/core/features/top_level_feature.py +++ b/tensorflow_datasets/core/features/top_level_feature.py @@ -180,7 +180,7 @@ def _example_parser_np(self): @functools.cached_property def _example_serializer(self): example_specs = self.get_serialized_info() - return example_serializer.ExampleSerializer(example_specs) + return example_serializer.ExampleSerializer(example_specs) # pyrefly: ignore[bad-argument-type] def _decode_feature( diff --git a/tensorflow_datasets/core/features/translation_feature.py b/tensorflow_datasets/core/features/translation_feature.py index c8b6323e57a..f9c1c3b0fd0 100644 --- a/tensorflow_datasets/core/features/translation_feature.py +++ b/tensorflow_datasets/core/features/translation_feature.py @@ -113,7 +113,7 @@ def languages(self): return sorted(self.keys()) @classmethod - def from_json_content( + def from_json_content( # pyrefly: ignore[bad-override] cls, value: Union[Json, feature_pb2.TranslationFeature] ) -> "Translation": if isinstance(value, dict): @@ -122,7 +122,7 @@ def from_json_content( "TFDS does not support datasets with Encoder. Please use the plain " "text version with `tensorflow_text`." ) - return cls(**value) + return cls(**value) # pyrefly: ignore[bad-argument-type] assert not value.variable_languages_per_example return cls(languages=value.languages) @@ -216,10 +216,10 @@ def languages(self): return sorted(list(self._languages)) if self._languages else None def encode_example(self, translation_dict): - if self.languages and set(translation_dict) - self._languages: + if self.languages and set(translation_dict) - self._languages: # pyrefly: ignore[unsupported-operation] raise ValueError( "Some languages in example ({0}) are not in valid set ({1}).".format( - ", ".join(sorted(set(translation_dict) - self._languages)), + ", ".join(sorted(set(translation_dict) - self._languages)), # pyrefly: ignore[unsupported-operation] ", ".join(self.languages), ) ) @@ -241,11 +241,11 @@ def encode_example(self, translation_dict): ) @classmethod - def from_json_content( + def from_json_content( # pyrefly: ignore[bad-override] cls, value: Union[Json, feature_pb2.TranslationFeature] ) -> "TranslationVariableLanguages": if isinstance(value, dict): - return cls(**value) + return cls(**value) # pyrefly: ignore[bad-argument-type] assert value.variable_languages_per_example return cls(languages=value.languages) diff --git a/tensorflow_datasets/core/features/video_feature.py b/tensorflow_datasets/core/features/video_feature.py index f5de49b25d7..f0c83264229 100644 --- a/tensorflow_datasets/core/features/video_feature.py +++ b/tensorflow_datasets/core/features/video_feature.py @@ -97,7 +97,7 @@ def __init__( encoding_format: str = 'png', ffmpeg_extra_args: Sequence[str] = (), use_colormap: bool = False, - dtype: Union[np.uint8, tf.uint8, np.uint16, tf.uint16] = np.uint8, + dtype: Union[np.uint8, tf.uint8, np.uint16, tf.uint16] = np.uint8, # pyrefly: ignore[not-a-type] doc: feature_lib.DocArg = None, ): """Initializes the connector. @@ -120,7 +120,7 @@ def __init__( Raises: ValueError: If the shape is invalid """ - dtype = tf.dtypes.as_dtype(dtype) + dtype = tf.dtypes.as_dtype(dtype) # pyrefly: ignore[bad-argument-type] frame_shape = None if shape: shape = tuple(shape) @@ -186,23 +186,23 @@ def encode_example(self, video_or_path_or_fobj): return super(Video, self).encode_example(encoded_video) @classmethod - def from_json_content( + def from_json_content( # pyrefly: ignore[bad-override] cls, value: Union[Json, feature_pb2.VideoFeature] ) -> 'Video': if isinstance(value, dict): # For backwards compatibility - shape = tuple(value['shape']) + shape = tuple(value['shape']) # pyrefly: ignore[bad-argument-type] encoding_format = value['encoding_format'] ffmpeg_extra_args = value['ffmpeg_extra_args'] return cls( shape=shape, - encoding_format=encoding_format, - ffmpeg_extra_args=ffmpeg_extra_args, + encoding_format=encoding_format, # pyrefly: ignore[bad-argument-type] + ffmpeg_extra_args=ffmpeg_extra_args, # pyrefly: ignore[bad-argument-type] ) return cls( shape=feature_lib.from_shape_proto(value.shape), dtype=feature_lib.dtype_from_str(value.dtype), - encoding_format=value.encoding_format or None, + encoding_format=value.encoding_format or None, # pyrefly: ignore[bad-argument-type] use_colormap=value.use_colormap, ffmpeg_extra_args=value.ffmpeg_extra_args, ) @@ -210,7 +210,7 @@ def from_json_content( def to_json_content(self) -> feature_pb2.VideoFeature: # pytype: disable=signature-mismatch # overriding-return-type-checks return feature_pb2.VideoFeature( 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] encoding_format=self._encoding_format, use_colormap=self._use_colormap, ffmpeg_extra_args=self._extra_ffmpeg_args, diff --git a/tensorflow_datasets/core/file_adapters.py b/tensorflow_datasets/core/file_adapters.py index ad3deb50a20..961102ef9c0 100644 --- a/tensorflow_datasets/core/file_adapters.py +++ b/tensorflow_datasets/core/file_adapters.py @@ -229,7 +229,7 @@ def make_tf_data( ) -> tf.data.Dataset: """Returns TensorFlow Dataset comprising given record file.""" buffer_size = buffer_size or cls.BUFFER_SIZE - return tf.data.TFRecordDataset(filename, buffer_size=buffer_size) + return tf.data.TFRecordDataset(filename, buffer_size=buffer_size) # pyrefly: ignore[bad-argument-type] @classmethod def write_examples( diff --git a/tensorflow_datasets/core/load.py b/tensorflow_datasets/core/load.py index ab25df6b594..b75c1992360 100644 --- a/tensorflow_datasets/core/load.py +++ b/tensorflow_datasets/core/load.py @@ -168,7 +168,7 @@ def builder( # 'kaggle:my_ds/config:1.0.0' -> ( # DatasetName('kaggle:my_ds'), {'version': '1.0.0', 'config': 'conf0'} # ) - name, builder_kwargs = naming.parse_builder_name_kwargs( + name, builder_kwargs = naming.parse_builder_name_kwargs( # pyrefly: ignore[bad-assignment] name, **builder_kwargs ) @@ -176,7 +176,7 @@ def get_dataset_repr() -> str: return f'dataset "{name}", builder_kwargs "{builder_kwargs}"' # `try_gcs` currently only supports non-community datasets - if try_gcs and not name.namespace and gcs_utils.is_dataset_on_gcs(str(name)): + if try_gcs and not name.namespace and gcs_utils.is_dataset_on_gcs(str(name)): # pyrefly: ignore[missing-attribute] data_dir = builder_kwargs.get('data_dir') if data_dir: raise ValueError( @@ -184,15 +184,15 @@ def get_dataset_repr() -> str: f' explicitly set. Wrong arguments for {get_dataset_repr()}' ) builder_kwargs['data_dir'] = gcs_utils.gcs_path('datasets') - if name.namespace: + if name.namespace: # pyrefly: ignore[missing-attribute] if name.namespace == 'huggingface': return huggingface_dataset_builder.builder( name=name.name, **builder_kwargs) if ( visibility.DatasetType.COMMUNITY_PUBLIC.is_available() - and community.community_register().has_namespace(name.namespace) + and community.community_register().has_namespace(name.namespace) # pyrefly: ignore[bad-argument-type] ): - return community.community_register().builder(name=name, **builder_kwargs) + return community.community_register().builder(name=name, **builder_kwargs) # pyrefly: ignore[bad-argument-type] # First check whether we can find the corresponding dataset builder code try: @@ -219,7 +219,7 @@ def get_dataset_repr() -> str: return cls(**builder_kwargs) # pytype: disable=not-instantiable # If neither the code nor the files are found, raise DatasetNotFoundError - raise not_found_error + raise not_found_error # pyrefly: ignore[unbound-name] def _try_load_from_files_first( @@ -343,7 +343,7 @@ def load_dataset( # If `split` is defined both as argument and in `loader_kwargs`, always keep # the one defined as argument. if split: - loader_kwargs['split'] = dataset_reference.get_split(split) + loader_kwargs['split'] = dataset_reference.get_split(split) # pyrefly: ignore[bad-argument-type] # Make sure we always return a dict of dicts. if 'split' in loader_kwargs and isinstance(loader_kwargs['split'], str): loader_kwargs['split'] = [loader_kwargs['split']] diff --git a/tensorflow_datasets/core/naming.py b/tensorflow_datasets/core/naming.py index d062194eca5..77748051691 100644 --- a/tensorflow_datasets/core/naming.py +++ b/tensorflow_datasets/core/naming.py @@ -434,18 +434,18 @@ def replace_func(matchobj): return f"{matchobj.group('prefix')}{replacement}" return f"{matchobj.group('prefix')}{replacement}{matchobj.group('suffix')}" - filepath = epath.Path(filepath) + filepath = epath.Path(filepath) # pyrefly: ignore[bad-assignment] (new_string, num_subs) = re.subn( pattern=pattern, repl=replace_func, - string=filepath.name, + string=filepath.name, # pyrefly: ignore[missing-attribute] ) if num_subs != 1: raise RuntimeError( f'Should do 1 shard suffix substitution, but did {num_subs}! ' f'Filepath was {filepath}' ) - return os.fspath(filepath.parent / new_string) + return os.fspath(filepath.parent / new_string) # pyrefly: ignore[missing-attribute] def _remove_extension(filepath: str) -> str: diff --git a/tensorflow_datasets/core/read_only_builder.py b/tensorflow_datasets/core/read_only_builder.py index b04ace6efa8..9b68ed9858e 100644 --- a/tensorflow_datasets/core/read_only_builder.py +++ b/tensorflow_datasets/core/read_only_builder.py @@ -88,10 +88,10 @@ def __init__( f' formats in the DatasetInfo: {sorted(available_formats)}.' ) - self.name = info_proto.name + self.name = info_proto.name # pyrefly: ignore[read-only] self.VERSION = version_lib.Version(info_proto.version) # pylint: disable=invalid-name - self.RELEASE_NOTES = info_proto.release_notes or {} # pylint: disable=invalid-name - self.BLOCKED_VERSIONS = self._restore_blocked_versions(info_proto) # pylint: disable=invalid-name + self.RELEASE_NOTES = info_proto.release_notes or {} # pylint: disable=invalid-name # pyrefly: ignore[read-only] + self.BLOCKED_VERSIONS = self._restore_blocked_versions(info_proto) # pylint: disable=invalid-name # pyrefly: ignore[read-only] if info_proto.module_name: # Overwrite the module so documenting `ReadOnlyBuilder` point to the @@ -135,7 +135,7 @@ def _restore_blocked_versions( configs = { info_proto.version: {info_proto.config_name: info_proto.is_blocked} } - return version_lib.BlockedVersions(configs=configs) + return version_lib.BlockedVersions(configs=configs) # pyrefly: ignore[bad-argument-type] return None def _create_builder_config( @@ -160,7 +160,7 @@ def _build_data_dir(self, data_dir: str) -> tuple[str, str]: def _info(self) -> dataset_info.DatasetInfo: return dataset_info.DatasetInfo.from_proto(self, self._info_proto) - def _download_and_prepare(self, **kwargs): # pylint: disable=arguments-differ + def _download_and_prepare(self, **kwargs): # pylint: disable=arguments-differ # pyrefly: ignore[bad-override] # DatasetBuilder.download_and_prepare is a no-op as self.data_dir already # exists. raise AssertionError("ReadOnlyBuilder can't be generated.") @@ -351,7 +351,7 @@ def _find_builder_dir(name: str, **builder_kwargs: Any) -> epath.Path | None: path: The dataset path found, or None if the dataset isn't found. """ # Normalize builder kwargs - name, builder_kwargs = naming.parse_builder_name_kwargs( + name, builder_kwargs = naming.parse_builder_name_kwargs( # pyrefly: ignore[bad-assignment] name, **builder_kwargs ) version = builder_kwargs.pop('version', None) @@ -366,13 +366,13 @@ def _find_builder_dir(name: str, **builder_kwargs: Any) -> epath.Path | None: # * config objects (rather than `str`) # * custom DatasetBuilder.__init__ kwargs if ( - name.namespace + name.namespace # pyrefly: ignore[missing-attribute] or version == 'experimental_latest' or isinstance(config, dataset_builder.BuilderConfig) or builder_kwargs ): error_msgs = ['Builder cannot be loaded from files if it uses:'] - if name.namespace: + if name.namespace: # pyrefly: ignore[missing-attribute] error_msgs.append(f'* namespaces (here, {name.namespace} is used)') if version == 'experimental_latest': error_msgs.append('* `experimental_latest` as requested version.') @@ -388,7 +388,7 @@ def _find_builder_dir(name: str, **builder_kwargs: Any) -> epath.Path | None: all_data_dirs = set(file_utils.list_data_dirs(given_data_dir=data_dir)) find_builder_fn = functools.partial( _find_builder_dir_single_dir, - builder_name=name.name, + builder_name=name.name, # pyrefly: ignore[missing-attribute] config_name=config, version=version, ) @@ -419,10 +419,11 @@ def wrapped_find_builder_fn(data_dir): # If the dataset root_dir exists, a common error is that the config name # was not specified. So we list the possible configs and display them. - possible_configs = _list_possible_configs(name.name, all_data_dirs) + possible_configs = _list_possible_configs(name.name, all_data_dirs) # pyrefly: ignore[missing-attribute] if possible_configs: configs_str = '\n\t- '.join([''] + possible_configs) error_msg = ( + # pyrefly: ignore[missing-attribute] f'However, a folder for "{name.name}" does exist. Is it possible that' ' you specified the wrong config? You can add a config by replacing' f' `tfds.load({name.name})` by `tfds.load("{name.name}/my_config")`.' @@ -555,7 +556,7 @@ def _get_default_config_name( else: # If code found, return the default config if cls.BUILDER_CONFIGS: - return cls.default_builder_config.name + return cls.default_builder_config.name # pyrefly: ignore[missing-attribute] # Otherwise, try to load default config from common metadata return dataset_builder.load_default_config_name(builder_dir) diff --git a/tensorflow_datasets/core/reader.py b/tensorflow_datasets/core/reader.py index 88a1445130f..d0b7b9ff36e 100644 --- a/tensorflow_datasets/core/reader.py +++ b/tensorflow_datasets/core/reader.py @@ -51,11 +51,11 @@ class _IdExample(NamedTuple): class _Instruction(NamedTuple): - filepath: tf.string # e.g. '/path/to/../train.tfrecord' - filename: tf.string # e.g. 'train.tfrecord' - tfds_id_prefix: tf.string # e.g. 'train.tfrecord' or 'folder1/train.tfrecord' - skip: tf.int64 - take: tf.int64 + filepath: tf.string # e.g. '/path/to/../train.tfrecord' # pyrefly: ignore[not-a-type] + filename: tf.string # e.g. 'train.tfrecord' # pyrefly: ignore[not-a-type] + tfds_id_prefix: tf.string # e.g. 'train.tfrecord' or 'folder1/train.tfrecord' # pyrefly: ignore[not-a-type] + skip: tf.int64 # pyrefly: ignore[not-a-type] + take: tf.int64 # pyrefly: ignore[not-a-type] def _get_dataset_from_filename( @@ -77,7 +77,7 @@ def _get_dataset_from_filename( if add_tfds_id: # For each example, generate a unique id. id_ds = _make_id_dataset( filename=instruction.tfds_id_prefix, - start_index=instruction.skip if do_skip else 0, + start_index=instruction.skip if do_skip else 0, # pyrefly: ignore[bad-argument-type] ) ds = tf.data.Dataset.zip(_IdExample(id=id_ds, example=ds)) return ds @@ -253,7 +253,7 @@ def validate_input_context(): if disable_shuffling: _verify_read_config_for_ordered_dataset( read_config, - interleave_cycle_length=cycle_length, + interleave_cycle_length=cycle_length, # pyrefly: ignore[bad-argument-type] shuffle_files=shuffle_files, ) @@ -282,7 +282,7 @@ def validate_input_context(): instruction_ds = instruction_ds.shuffle( len(file_instructions), seed=read_config.shuffle_seed, - reshuffle_each_iteration=read_config.shuffle_reshuffle_each_iteration, + reshuffle_each_iteration=read_config.shuffle_reshuffle_each_iteration, # pyrefly: ignore[bad-argument-type] ) if read_config.repeat_filenames: @@ -309,7 +309,7 @@ def validate_input_context(): add_tfds_id=read_config.add_tfds_id, override_buffer_size=read_config.override_buffer_size, ), - cycle_length=cycle_length, + cycle_length=cycle_length, # pyrefly: ignore[bad-argument-type] block_length=block_length, num_parallel_calls=read_config.num_parallel_calls_for_interleave_files, deterministic=deterministic, @@ -478,7 +478,7 @@ def parse_and_decode(ex: Tensor) -> TreeDict[Tensor]: # https://www.tensorflow.org/api_docs/python/tf/io/parse_example ex = self._parser.parse_example(ex) if decode_fn: - ex = decode_fn(ex) + ex = decode_fn(ex) # pyrefly: ignore[bad-assignment] return ex # Eventually add the `tfds_id` after the decoding diff --git a/tensorflow_datasets/core/sequential_writer.py b/tensorflow_datasets/core/sequential_writer.py index 34fae4f450f..79de90c3722 100644 --- a/tensorflow_datasets/core/sequential_writer.py +++ b/tensorflow_datasets/core/sequential_writer.py @@ -58,7 +58,7 @@ class Shard(object): def add_example(self, serialized_example: str) -> None: """Writes a new example.""" - self.writer.write(serialized_example) + self.writer.write(serialized_example) # pyrefly: ignore[bad-argument-type] self.num_examples += 1 self.num_bytes += len(serialized_example) diff --git a/tensorflow_datasets/core/shuffle.py b/tensorflow_datasets/core/shuffle.py index 541fdf15a19..b005fcecbf7 100644 --- a/tensorflow_datasets/core/shuffle.py +++ b/tensorflow_datasets/core/shuffle.py @@ -161,12 +161,12 @@ def add(self, key: type_utils.Key, data: bytes): data_size = len(data) try: - self._fobj.write(_hkey_to_bytes(key)) + self._fobj.write(_hkey_to_bytes(key)) # pyrefly: ignore[bad-argument-type] except tf.errors.ResourceExhaustedError as error: # catch "Too many open files" if error.message.endswith('Too many open files'): _increase_open_files_limit() - self._fobj.write(_hkey_to_bytes(key)) + self._fobj.write(_hkey_to_bytes(key)) # pyrefly: ignore[bad-argument-type] else: raise error # http://docs.python.org/3/library/struct.html#byte-order-size-and-alignment @@ -272,10 +272,10 @@ def _add_to_mem_buffer(self, hkey: int, data: bytes) -> None: if self._total_bytes > MAX_MEM_BUFFER_SIZE: for hkey, data in self._mem_buffer: self._add_to_bucket(hkey, data) - self._mem_buffer = None + self._mem_buffer = None # pyrefly: ignore[bad-assignment] self._in_memory = False - def add(self, key: type_utils.Key, data: bytes) -> bool: + def add(self, key: type_utils.Key, data: bytes) -> bool: # pyrefly: ignore[bad-return] """Add (key, data) to shuffler.""" if self._read_only: raise AssertionError('add() cannot be called after __iter__.') diff --git a/tensorflow_datasets/core/split_builder.py b/tensorflow_datasets/core/split_builder.py index 5fa705078ab..358e757b64d 100644 --- a/tensorflow_datasets/core/split_builder.py +++ b/tensorflow_datasets/core/split_builder.py @@ -181,7 +181,7 @@ def submit_shard_based_generation( num_shards = len(example_gen_per_shard) filename_template = filename_template.replace(split=split_name) serialized_info = self._features.get_serialized_info() - serializer = example_serializer.ExampleSerializer(serialized_info) + serializer = example_serializer.ExampleSerializer(serialized_info) # pyrefly: ignore[bad-argument-type] shard_writer = writer_lib.ShardWriter( features=self._features, @@ -198,7 +198,7 @@ def submit_shard_based_generation( ) shard_paths.append(shard_path) num_examples = shard_writer.write( - path=shard_path, examples=example_gen() + path=shard_path, examples=example_gen() # pyrefly: ignore[bad-argument-type] ) shard_lengths.append(num_examples) else: @@ -217,7 +217,7 @@ def submit_shard_based_generation( shard_paths.append(shard_path) shard_info = shard_writer.write_with_beam( path=shard_path, - example_gen=example_gen, + example_gen=example_gen, # pyrefly: ignore[bad-argument-type] shard_index=shard_index, pipeline=self.beam_pipeline, ) @@ -457,9 +457,9 @@ def submit_split_generation( # Generate the beam.PCollection pcollection = self.beam_pipeline | split_name >> generator build_kwargs['generator'] = pcollection - return self._build_from_pcollection(**build_kwargs) + return self._build_from_pcollection(**build_kwargs) # pyrefly: ignore[missing-argument] elif isinstance(generator, beam.PCollection): - return self._build_from_pcollection(**build_kwargs) + return self._build_from_pcollection(**build_kwargs) # pyrefly: ignore[missing-argument] else: raise unknown_generator_type @@ -498,7 +498,7 @@ def _build_from_generator( serialized_info = self._features.get_serialized_info() writer = writer_lib.Writer( - serializer=example_serializer.ExampleSerializer(serialized_info), + serializer=example_serializer.ExampleSerializer(serialized_info), # pyrefly: ignore[bad-argument-type] filename_template=filename_template, hash_salt=split_name, disable_shuffling=disable_shuffling, @@ -550,7 +550,7 @@ def _build_from_pcollection( """Split generator for `beam.PCollection`.""" # TODO(tfds): Should try to add support to `max_examples_per_split` serializer = example_serializer.ExampleSerializer( - self._features.get_serialized_info() + self._features.get_serialized_info() # pyrefly: ignore[bad-argument-type] ) if nondeterministic_order: logging.info( @@ -560,7 +560,7 @@ def _build_from_pcollection( beam_writer = writer_lib.NoShuffleBeamWriter( serializer=serializer, file_format=file_adapters.FileFormat.from_value( - filename_template.filetype_suffix + filename_template.filetype_suffix # pyrefly: ignore[bad-argument-type] ), filename_template=filename_template, num_shards=num_shards, diff --git a/tensorflow_datasets/core/splits.py b/tensorflow_datasets/core/splits.py index 0717f113f44..49028a1c768 100644 --- a/tensorflow_datasets/core/splits.py +++ b/tensorflow_datasets/core/splits.py @@ -444,7 +444,7 @@ def __repr__(self) -> str: if typing.TYPE_CHECKING: # For type checking, `tfds.Split` is an alias for `str` with additional # `.TRAIN`, `.TEST`,... attributes. All strings are valid split type. - Split = Split | str + Split = Split | str # pyrefly: ignore[bad-assignment] class SplitDict(utils.NonMutableDict[str, SplitInfo]): @@ -468,7 +468,7 @@ def __init__( ) self._dataset_name = dataset_name # deprecated, please don't use - def __getitem__(self, key) -> SplitInfo | SubSplitInfo: + def __getitem__(self, key) -> SplitInfo | SubSplitInfo: # pyrefly: ignore[bad-override] if not self: raise KeyError( f'Trying to access `splits[{key!r}]` but `splits` is empty. ' @@ -797,16 +797,16 @@ def _str_to_relative_instruction(spec: str) -> AbstractSplit: raise ValueError(err_msg) if len(slices) == 1: # split='train[x]' (from_match,) = slices - from_ = from_match['val'] + from_ = from_match['val'] # pyrefly: ignore[unsupported-operation] to = int(from_) + 1 - unit = from_match['unit'] or 'abs' + unit = from_match['unit'] or 'abs' # pyrefly: ignore[unsupported-operation] if unit != 'shard': raise ValueError('Absolute or percent only support slice syntax.') elif len(slices) == 2: # split='train[x:y]' from_match, to_match = slices - from_ = from_match['val'] - to = to_match['val'] - unit = from_match['unit'] or to_match['unit'] or 'abs' + from_ = from_match['val'] # pyrefly: ignore[unsupported-operation] + to = to_match['val'] # pyrefly: ignore[unsupported-operation] + unit = from_match['unit'] or to_match['unit'] or 'abs' # pyrefly: ignore[unsupported-operation] else: raise ValueError(err_msg) @@ -867,10 +867,10 @@ def _rel_to_abs_instr( to = num_examples if to is None else pct_to_abs(to, num_examples) elif rel_instr.unit == 'shard': shard_lengths = split_infos[split].shard_lengths - from_ = 0 if from_ is None else sum(shard_lengths[:from_]) + from_ = 0 if from_ is None else sum(shard_lengths[:from_]) # pyrefly: ignore[bad-index] if to is not None and to <= 0: to = len(shard_lengths) + to - to = num_examples if to is None else sum(shard_lengths[:to]) + to = num_examples if to is None else sum(shard_lengths[:to]) # pyrefly: ignore[bad-index] elif rel_instr.unit == 'abs': from_ = 0 if from_ is None else from_ to = num_examples if to is None else to @@ -891,4 +891,4 @@ def _rel_to_abs_instr( to = num_examples + to elif to == num_examples: to = None - return _AbsoluteInstruction(split, from_, to) + return _AbsoluteInstruction(split, from_, to) # pyrefly: ignore[bad-argument-type] diff --git a/tensorflow_datasets/core/subsplits_utils.py b/tensorflow_datasets/core/subsplits_utils.py index e82b0ce4777..c651b750186 100644 --- a/tensorflow_datasets/core/subsplits_utils.py +++ b/tensorflow_datasets/core/subsplits_utils.py @@ -46,7 +46,7 @@ def to_absolute(self, split_infos: splits_lib.SplitDict): self._absolute_to_read_instruction_for_index(abs_inst, split_infos) for abs_inst in absolute_instructions ] - instuctions = functools.reduce(operator.add, read_instructions_for_index) + instuctions = functools.reduce(operator.add, read_instructions_for_index) # pyrefly: ignore[incompatible-overload-residual] return instuctions.to_absolute(split_infos) def _absolute_to_read_instruction_for_index( diff --git a/tensorflow_datasets/core/writer.py b/tensorflow_datasets/core/writer.py index 080d8be3dff..8b8f8eacaa1 100644 --- a/tensorflow_datasets/core/writer.py +++ b/tensorflow_datasets/core/writer.py @@ -147,7 +147,7 @@ def _get_shard_specs( _ShardSpec( shard_index=shard_index, path=os.fspath(shard_path), - index_path=index_path, + index_path=index_path, # pyrefly: ignore[bad-argument-type] examples_number=to - from_, file_instructions=file_instructions, ) @@ -620,9 +620,9 @@ def _assign_shard( ) -> tuple[int, tuple[Any, bytes]]: """Assigns a shard id to the example.""" key, _ = key_serialized_example - largest_key = largest_key[0] + largest_key = largest_key[0] # pyrefly: ignore[bad-assignment] shard_number = shuffle.get_bucket_number( - hkey=key, num_buckets=num_shards, max_hkey=largest_key + hkey=key, num_buckets=num_shards, max_hkey=largest_key # pyrefly: ignore[bad-argument-type] ) self._get_distribution(name="ShardDistribution").update(shard_number) self.inc_counter(f"{self._filename_template.split}.shard_{shard_number}")