From 1bf64acdb0fa5c5ed7f2ebe281810fa15528aa98 Mon Sep 17 00:00:00 2001 From: Gautam Kishore Date: Tue, 14 Jul 2026 15:45:47 +0530 Subject: [PATCH] fix: replace list/List with Sequence in function parameter annotations Changes list[T] to Sequence[T] in function parameter type annotations across arrow_dataset.py, dataset_dict.py, and iterable_dataset.py. Sequence is covariant, so passing List[str] where Sequence[PathLike] is expected no longer triggers a mypy error. Fixes #5354 --- src/datasets/arrow_dataset.py | 62 ++++++++++++++++---------------- src/datasets/dataset_dict.py | 22 ++++++------ src/datasets/iterable_dataset.py | 62 ++++++++++++++++---------------- 3 files changed, 73 insertions(+), 73 deletions(-) diff --git a/src/datasets/arrow_dataset.py b/src/datasets/arrow_dataset.py index 59451a640e6..e7cc156a04b 100644 --- a/src/datasets/arrow_dataset.py +++ b/src/datasets/arrow_dataset.py @@ -249,7 +249,7 @@ def _get_output_signature( dataset: "Dataset", collate_fn: Callable, collate_fn_args: dict, - cols_to_retain: Optional[list[str]] = None, + cols_to_retain: Optional[Sequence_[str]] = None, batch_size: Optional[int] = None, num_test_batches: int = 20, ): @@ -345,12 +345,12 @@ def _get_output_signature( def to_tf_dataset( self, batch_size: Optional[int] = None, - columns: Optional[Union[str, list[str]]] = None, + columns: Optional[Union[str, Sequence_[str]]] = None, shuffle: bool = False, collate_fn: Optional[Callable] = None, drop_remainder: bool = False, collate_fn_args: Optional[dict[str, Any]] = None, - label_cols: Optional[Union[str, list[str]]] = None, + label_cols: Optional[Union[str, Sequence_[str]]] = None, prefetch: bool = True, num_workers: int = 0, num_test_batches: int = 20, @@ -625,7 +625,7 @@ def _check_table(table) -> Table: raise TypeError(f"Expected a pyarrow.Table or a datasets.table.Table object, but got {table}.") -def _check_column_names(column_names: list[str]): +def _check_column_names(column_names: Sequence_[str]): """Check the column names to make sure they don't contain duplicates.""" counter = Counter(column_names) if not all(count == 1 for count in counter.values()): @@ -684,7 +684,7 @@ def __iter__(self) -> Iterator[Any]: for example in source: yield example[self.column_name] - def __getitem__(self, key: Union[int, str, list[int]]) -> Any: + def __getitem__(self, key: Union[int, str, Sequence_[int]]) -> Any: if isinstance(key, str): return Column(self, key) elif isinstance(self.source, Dataset): @@ -1157,7 +1157,7 @@ def from_dict( @classmethod def from_list( cls, - mapping: list[dict], + mapping: Sequence_[dict], features: Optional[Features] = None, info: Optional[DatasetInfo] = None, split: Optional[NamedSplit] = None, @@ -1290,7 +1290,7 @@ def from_list( @staticmethod def from_csv( - path_or_paths: Union[PathLike, list[PathLike]], + path_or_paths: Union[PathLike, Sequence_[PathLike]], split: Optional[NamedSplit] = None, features: Optional[Features] = None, cache_dir: str = None, @@ -1430,7 +1430,7 @@ def from_generator( @staticmethod def from_json( - path_or_paths: Union[PathLike, list[PathLike]], + path_or_paths: Union[PathLike, Sequence_[PathLike]], split: Optional[NamedSplit] = None, features: Optional[Features] = None, cache_dir: str = None, @@ -1489,14 +1489,14 @@ def from_json( @staticmethod def from_parquet( - path_or_paths: Union[PathLike, list[PathLike]], + path_or_paths: Union[PathLike, Sequence_[PathLike]], split: Optional[NamedSplit] = None, features: Optional[Features] = None, cache_dir: str = None, keep_in_memory: bool = False, - columns: Optional[list[str]] = None, + columns: Optional[Sequence_[str]] = None, num_proc: Optional[int] = None, - filters: Optional[Union[pds.Expression, list[tuple], list[list[tuple]]]] = None, + filters: Optional[Union[pds.Expression, Sequence_[tuple], Sequence_[Sequence_[tuple]]]] = None, fragment_scan_options: Optional[pds.ParquetFragmentScanOptions] = None, on_bad_files: Literal["error", "warn", "skip"] = "error", **kwargs, @@ -1586,7 +1586,7 @@ def from_parquet( @staticmethod def from_text( - path_or_paths: Union[PathLike, list[PathLike]], + path_or_paths: Union[PathLike, Sequence_[PathLike]], split: Optional[NamedSplit] = None, features: Optional[Features] = None, cache_dir: str = None, @@ -2488,7 +2488,7 @@ def cast_column(self, column: str, feature: FeatureType, new_fingerprint: Option @transmit_format @fingerprint_transform(inplace=False) - def remove_columns(self, column_names: Union[str, list[str]], new_fingerprint: Optional[str] = None) -> "Dataset": + def remove_columns(self, column_names: Union[str, Sequence_[str]], new_fingerprint: Optional[str] = None) -> "Dataset": """ Remove one or several column(s) in the dataset and the features associated to them. @@ -2676,7 +2676,7 @@ def rename(columns): @transmit_format @fingerprint_transform(inplace=False) - def select_columns(self, column_names: Union[str, list[str]], new_fingerprint: Optional[str] = None) -> "Dataset": + def select_columns(self, column_names: Union[str, Sequence_[str]], new_fingerprint: Optional[str] = None) -> "Dataset": """Select one or several column(s) in the dataset and the features associated to them. @@ -3216,11 +3216,11 @@ def map( function: Optional[Callable] = None, with_indices: bool = False, with_rank: bool = False, - input_columns: Optional[Union[str, list[str]]] = None, + input_columns: Optional[Union[str, Sequence_[str]]] = None, batched: bool = False, batch_size: Optional[int] = 1000, drop_last_batch: bool = False, - remove_columns: Optional[Union[str, list[str]]] = None, + remove_columns: Optional[Union[str, Sequence_[str]]] = None, keep_in_memory: bool = False, load_from_cache_file: Optional[bool] = None, cache_file_name: Optional[str] = None, @@ -3671,11 +3671,11 @@ def _map_single( function: Optional[Callable] = None, with_indices: bool = False, with_rank: bool = False, - input_columns: Optional[list[str]] = None, + input_columns: Optional[Sequence_[str]] = None, batched: bool = False, batch_size: Optional[int] = 1000, drop_last_batch: bool = False, - remove_columns: Optional[list[str]] = None, + remove_columns: Optional[Sequence_[str]] = None, keep_in_memory: bool = False, cache_file_name: Optional[str] = None, writer_batch_size: Optional[int] = 1000, @@ -4064,7 +4064,7 @@ def iter_outputs(shard_iterable): def batch( self, batch_size: Optional[int] = None, - by_column: Optional[Union[str, list[str]]] = None, + by_column: Optional[Union[str, Sequence_[str]]] = None, drop_last_batch: bool = False, num_proc: Optional[int] = None, new_fingerprint: Optional[str] = None, @@ -4145,7 +4145,7 @@ def filter( function: Optional[Callable] = None, with_indices: bool = False, with_rank: bool = False, - input_columns: Optional[Union[str, list[str]]] = None, + input_columns: Optional[Union[str, Sequence_[str]]] = None, batched: bool = False, batch_size: Optional[int] = 1000, keep_in_memory: bool = False, @@ -4717,9 +4717,9 @@ def sort( """Create a new dataset sorted according to a single or multiple columns. Args: - column_names (`Union[str, Sequence[str]]`): + column_names (`Union[str, Sequence_[str]]`): Column name(s) to sort by. - reverse (`Union[bool, Sequence[bool]]`, defaults to `False`): + reverse (`Union[bool, Sequence_[bool]]`, defaults to `False`): If `True`, sort by descending order rather than ascending. If a single bool is provided, the value is applied to the sorting of all column names. Otherwise a list of bools with the same length and order as column_names must be provided. @@ -5725,7 +5725,7 @@ def extra_nbytes_visitor(array, feature): return dataset_nbytes @staticmethod - def _generate_tables_from_shards(shards: list["Dataset"], batch_size: int): + def _generate_tables_from_shards(shards: Sequence_["Dataset"], batch_size: int): for shard_idx, shard in enumerate(shards): for pa_table in shard.with_format("arrow").iter(batch_size): yield shard_idx, pa_table @@ -6872,12 +6872,12 @@ def _push_to_bucket( def _get_updated_dataset_card( fs: DirFileSystem, config_name: str, - splits_info: list[SplitInfo], + splits_info: Sequence_[SplitInfo], features: Features, data_dir: str, set_default: Optional[bool], - uploaded_sizes: list[int], - deleted_sizes: list[int], + uploaded_sizes: Sequence_[int], + deleted_sizes: Sequence_[int], remove_other_splits: bool, ) -> tuple[DatasetCard, Optional[dict]]: """Update a dataset card in push_to_hub""" @@ -7003,7 +7003,7 @@ def _get_updated_dataset_card( def _concatenate_map_style_datasets( - dsets: list[Dataset], + dsets: Sequence_[Dataset], info: Optional[DatasetInfo] = None, split: Optional[NamedSplit] = None, axis: int = 0, @@ -7117,8 +7117,8 @@ def apply_offset_to_indices_table(table, offset): def _interleave_map_style_datasets( - datasets: list["Dataset"], - probabilities: Optional[list[float]] = None, + datasets: Sequence_["Dataset"], + probabilities: Optional[Sequence_[float]] = None, seed: Optional[int] = None, info: Optional[DatasetInfo] = None, split: Optional[NamedSplit] = None, @@ -7279,7 +7279,7 @@ def get_indices_from_mask_function( batched: bool, with_indices: bool, with_rank: bool, - input_columns: Optional[Union[str, list[str]]], + input_columns: Optional[Union[str, Sequence_[str]]], indices_mapping: Optional[Table] = None, *args, **fn_kwargs, @@ -7337,7 +7337,7 @@ async def async_get_indices_from_mask_function( batched: bool, with_indices: bool, with_rank: bool, - input_columns: Optional[Union[str, list[str]]], + input_columns: Optional[Union[str, Sequence_[str]]], indices_mapping: Optional[Table] = None, *args, **fn_kwargs, diff --git a/src/datasets/dataset_dict.py b/src/datasets/dataset_dict.py index 4abea0a381a..cd1dc738928 100644 --- a/src/datasets/dataset_dict.py +++ b/src/datasets/dataset_dict.py @@ -342,7 +342,7 @@ def cast_column(self, column: str, feature) -> "DatasetDict": self._check_values_type() return DatasetDict({k: dataset.cast_column(column=column, feature=feature) for k, dataset in self.items()}) - def remove_columns(self, column_names: Union[str, list[str]]) -> "DatasetDict": + def remove_columns(self, column_names: Union[str, Sequence[str]]) -> "DatasetDict": """ Remove one or several column(s) from each split in the dataset and the features associated to the column(s). @@ -470,7 +470,7 @@ def rename_columns(self, column_mapping: dict[str, str]) -> "DatasetDict": self._check_values_type() return DatasetDict({k: dataset.rename_columns(column_mapping=column_mapping) for k, dataset in self.items()}) - def select_columns(self, column_names: Union[str, list[str]]) -> "DatasetDict": + def select_columns(self, column_names: Union[str, Sequence[str]]) -> "DatasetDict": """Select one or several column(s) from each split in the dataset and the features associated to the column(s). @@ -827,11 +827,11 @@ def map( with_indices: bool = False, with_rank: bool = False, with_split: bool = False, - input_columns: Optional[Union[str, list[str]]] = None, + input_columns: Optional[Union[str, Sequence[str]]] = None, batched: bool = False, batch_size: Optional[int] = 1000, drop_last_batch: bool = False, - remove_columns: Optional[Union[str, list[str]]] = None, + remove_columns: Optional[Union[str, Sequence[str]]] = None, keep_in_memory: bool = False, load_from_cache_file: Optional[bool] = None, cache_file_names: Optional[dict[str, Optional[str]]] = None, @@ -998,7 +998,7 @@ def filter( function: Optional[Callable] = None, with_indices: bool = False, with_rank: bool = False, - input_columns: Optional[Union[str, list[str]]] = None, + input_columns: Optional[Union[str, Sequence[str]]] = None, batched: bool = False, batch_size: Optional[int] = 1000, keep_in_memory: bool = False, @@ -1534,7 +1534,7 @@ def from_parquet( features: Optional[Features] = None, cache_dir: str = None, keep_in_memory: bool = False, - columns: Optional[list[str]] = None, + columns: Optional[Sequence[str]] = None, **kwargs, ) -> "DatasetDict": """Create [`DatasetDict`] from Parquet file(s). @@ -1925,11 +1925,11 @@ def map( function: Optional[Callable] = None, with_indices: bool = False, with_split: bool = False, - input_columns: Optional[Union[str, list[str]]] = None, + input_columns: Optional[Union[str, Sequence[str]]] = None, batched: bool = False, batch_size: int = 1000, drop_last_batch: bool = False, - remove_columns: Optional[Union[str, list[str]]] = None, + remove_columns: Optional[Union[str, Sequence[str]]] = None, fn_kwargs: Optional[dict] = None, ) -> "IterableDatasetDict": """ @@ -2024,7 +2024,7 @@ def filter( self, function: Optional[Callable] = None, with_indices=False, - input_columns: Optional[Union[str, list[str]]] = None, + input_columns: Optional[Union[str, Sequence[str]]] = None, batched: bool = False, batch_size: Optional[int] = 1000, fn_kwargs: Optional[dict] = None, @@ -2208,7 +2208,7 @@ def rename_columns(self, column_mapping: dict[str, str]) -> "IterableDatasetDict {k: dataset.rename_columns(column_mapping=column_mapping) for k, dataset in self.items()} ) - def remove_columns(self, column_names: Union[str, list[str]]) -> "IterableDatasetDict": + def remove_columns(self, column_names: Union[str, Sequence[str]]) -> "IterableDatasetDict": """ Remove one or several column(s) in the dataset and the features associated to them. The removal is done on-the-fly on the examples when iterating over the dataset. @@ -2234,7 +2234,7 @@ def remove_columns(self, column_names: Union[str, list[str]]) -> "IterableDatase """ return IterableDatasetDict({k: dataset.remove_columns(column_names) for k, dataset in self.items()}) - def select_columns(self, column_names: Union[str, list[str]]) -> "IterableDatasetDict": + def select_columns(self, column_names: Union[str, Sequence[str]]) -> "IterableDatasetDict": """Select one or several column(s) in the dataset and the features associated to them. The selection is done on-the-fly on the examples when iterating over the dataset. The selection is applied to all the diff --git a/src/datasets/iterable_dataset.py b/src/datasets/iterable_dataset.py index 17b9a2020fc..46ef3c26d50 100644 --- a/src/datasets/iterable_dataset.py +++ b/src/datasets/iterable_dataset.py @@ -9,7 +9,7 @@ import tempfile import time from collections import Counter -from collections.abc import Iterable, Iterator +from collections.abc import Iterable, Iterator, Sequence from copy import copy, deepcopy from dataclasses import dataclass from functools import partial @@ -120,7 +120,7 @@ def _rename_columns_fn(example: dict, column_mapping: dict[str, str]): } -def add_column_fn(example: dict, idx: int, name: str, column: list[dict]): +def add_column_fn(example: dict, idx: int, name: str, column: Sequence[dict]): if name in example: raise ValueError(f"Error when adding {name}: column {name} is already in the dataset.") return {name: column[idx]} @@ -136,7 +136,7 @@ def _infer_features_from_batch(batch: dict[str, list], try_features: Optional[Fe return Features.from_arrow_schema(pa_table.schema) -def _examples_to_batch(examples: list[dict[str, Any]]) -> dict[str, list]: +def _examples_to_batch(examples: Sequence[dict[str, Any]]) -> dict[str, list]: # we order the columns by order of appearance # to do so, we use a dict as an ordered set cols = {col: None for example in examples for col in example} @@ -622,7 +622,7 @@ def num_shards(self) -> int: class SelectColumnsIterable(_BaseExamplesIterable): - def __init__(self, ex_iterable: _BaseExamplesIterable, column_names: list[str]): + def __init__(self, ex_iterable: _BaseExamplesIterable, column_names: Sequence[str]): super().__init__() self.ex_iterable = ex_iterable self.column_names = column_names @@ -743,7 +743,7 @@ def num_shards(self) -> int: class CyclingMultiSourcesExamplesIterable(_BaseExamplesIterable): def __init__( self, - ex_iterables: list[_BaseExamplesIterable], + ex_iterables: Sequence[_BaseExamplesIterable], stopping_strategy: Literal[ "first_exhausted", "all_exhausted", "all_exhausted_without_replacement" ] = "first_exhausted", @@ -998,7 +998,7 @@ class VerticallyConcatenatedMultiSourcesExamplesIterable(_BaseExamplesIterable): This is done with `_apply_feature_types_on_example`. """ - def __init__(self, ex_iterables: list[_BaseExamplesIterable]): + def __init__(self, ex_iterables: Sequence[_BaseExamplesIterable]): super().__init__() self.ex_iterables = ex_iterables @@ -1074,7 +1074,7 @@ def reshard_data_sources(self) -> "VerticallyConcatenatedMultiSourcesExamplesIte ) -def _check_column_names(column_names: list[str]): +def _check_column_names(column_names: Sequence[str]): """Check the column names to make sure they don't contain duplicates.""" counter = Counter(column_names) if not all(count == 1 for count in counter.values()): @@ -1100,7 +1100,7 @@ class HorizontallyConcatenatedMultiSourcesExamplesIterable(_BaseExamplesIterable This is done with `_apply_feature_types_on_example`. """ - def __init__(self, ex_iterables: list[_BaseExamplesIterable]): + def __init__(self, ex_iterables: Sequence[_BaseExamplesIterable]): super().__init__() self.ex_iterables = ex_iterables @@ -1206,9 +1206,9 @@ def reshard_data_sources(self) -> "HorizontallyConcatenatedMultiSourcesExamplesI class RandomlyCyclingMultiSourcesExamplesIterable(CyclingMultiSourcesExamplesIterable): def __init__( self, - ex_iterables: list[_BaseExamplesIterable], + ex_iterables: Sequence[_BaseExamplesIterable], generator: np.random.Generator, - probabilities: Optional[list[float]] = None, + probabilities: Optional[Sequence[float]] = None, stopping_strategy: Literal[ "first_exhausted", "all_exhausted", "all_exhausted_without_replacement" ] = "first_exhausted", @@ -1347,11 +1347,11 @@ def __init__( ex_iterable: _BaseExamplesIterable, function: Callable, with_indices: bool = False, - input_columns: Optional[list[str]] = None, + input_columns: Optional[Sequence[str]] = None, batched: bool = False, batch_size: Optional[int] = 1000, drop_last_batch: bool = False, - remove_columns: Optional[list[str]] = None, + remove_columns: Optional[Sequence[str]] = None, fn_kwargs: Optional[dict] = None, formatting: Optional["FormattingConfig"] = None, features: Optional[Features] = None, @@ -1799,7 +1799,7 @@ def __init__( ex_iterable: _BaseExamplesIterable, function: Callable, with_indices: bool = False, - input_columns: Optional[list[str]] = None, + input_columns: Optional[Sequence[str]] = None, batched: bool = False, batch_size: Optional[int] = 1000, fn_kwargs: Optional[dict] = None, @@ -3123,7 +3123,7 @@ def from_dict( @classmethod def from_list( cls, - mapping: list[dict], + mapping: Sequence[dict], features: Optional[Features] = None, info: Optional[DatasetInfo] = None, split: Optional[NamedSplit] = None, @@ -3161,7 +3161,7 @@ def from_list( @staticmethod def from_csv( - path_or_paths: Union[PathLike, list[PathLike]], + path_or_paths: Union[PathLike, Sequence[PathLike]], split: Optional[NamedSplit] = None, features: Optional[Features] = None, keep_in_memory: bool = False, @@ -3204,7 +3204,7 @@ def from_csv( @staticmethod def from_json( - path_or_paths: Union[PathLike, list[PathLike]], + path_or_paths: Union[PathLike, Sequence[PathLike]], split: Optional[NamedSplit] = None, features: Optional[Features] = None, keep_in_memory: bool = False, @@ -3251,12 +3251,12 @@ def from_json( @staticmethod def from_parquet( - path_or_paths: Union[PathLike, list[PathLike]], + path_or_paths: Union[PathLike, Sequence[PathLike]], split: Optional[NamedSplit] = None, features: Optional[Features] = None, keep_in_memory: bool = False, - columns: Optional[list[str]] = None, - filters: Optional[Union[pds.Expression, list[tuple], list[list[tuple]]]] = None, + columns: Optional[Sequence[str]] = None, + filters: Optional[Union[pds.Expression, Sequence[tuple], Sequence[Sequence[tuple]]]] = None, fragment_scan_options: Optional[pds.ParquetFragmentScanOptions] = None, on_bad_files: Literal["error", "warn", "skip"] = "error", **kwargs, @@ -3336,7 +3336,7 @@ def from_parquet( @staticmethod def from_text( - path_or_paths: Union[PathLike, list[PathLike]], + path_or_paths: Union[PathLike, Sequence[PathLike]], split: Optional[NamedSplit] = None, features: Optional[Features] = None, keep_in_memory: bool = False, @@ -3447,11 +3447,11 @@ def map( self, function: Optional[Callable] = None, with_indices: bool = False, - input_columns: Optional[Union[str, list[str]]] = None, + input_columns: Optional[Union[str, Sequence[str]]] = None, batched: bool = False, batch_size: Optional[int] = 1000, drop_last_batch: bool = False, - remove_columns: Optional[Union[str, list[str]]] = None, + remove_columns: Optional[Union[str, Sequence[str]]] = None, features: Optional[Features] = None, fn_kwargs: Optional[dict] = None, ) -> "IterableDataset": @@ -3542,11 +3542,11 @@ def _map( self, function: Optional[Callable] = None, with_indices: bool = False, - input_columns: Optional[Union[str, list[str]]] = None, + input_columns: Optional[Union[str, Sequence[str]]] = None, batched: bool = False, batch_size: Optional[int] = 1000, drop_last_batch: bool = False, - remove_columns: Optional[Union[str, list[str]]] = None, + remove_columns: Optional[Union[str, Sequence[str]]] = None, features: Optional[Features] = None, fn_kwargs: Optional[dict] = None, is_batch_accumulate_arrow_table_function: bool = False, @@ -3629,7 +3629,7 @@ def filter( self, function: Optional[Callable] = None, with_indices=False, - input_columns: Optional[Union[str, list[str]]] = None, + input_columns: Optional[Union[str, Sequence[str]]] = None, batched: bool = False, batch_size: Optional[int] = 1000, fn_kwargs: Optional[dict] = None, @@ -4093,7 +4093,7 @@ def rename_columns(self, column_mapping: dict[str, str]) -> "IterableDataset": ) return ds_iterable - def remove_columns(self, column_names: Union[str, list[str]]) -> "IterableDataset": + def remove_columns(self, column_names: Union[str, Sequence[str]]) -> "IterableDataset": """ Remove one or several column(s) in the dataset and the features associated to them. The removal is done on-the-fly on the examples when iterating over the dataset. @@ -4128,7 +4128,7 @@ def remove_columns(self, column_names: Union[str, list[str]]) -> "IterableDatase return ds_iterable - def select_columns(self, column_names: Union[str, list[str]]) -> "IterableDataset": + def select_columns(self, column_names: Union[str, Sequence[str]]) -> "IterableDataset": """Select one or several column(s) in the dataset and the features associated to them. The selection is done on-the-fly on the examples when iterating over the dataset. @@ -4395,7 +4395,7 @@ def _resolve_features(self): def batch( self, batch_size: Optional[int] = None, - by_column: Optional[Union[str, list[str]]] = None, + by_column: Optional[Union[str, Sequence[str]]] = None, drop_last_batch: bool = False, ) -> "IterableDataset": """ @@ -5169,7 +5169,7 @@ def push_to_hub( def _concatenate_iterable_datasets( - dsets: list[IterableDataset], + dsets: Sequence[IterableDataset], info: Optional[DatasetInfo] = None, split: Optional[NamedSplit] = None, axis: int = 0, @@ -5261,8 +5261,8 @@ def _concatenate_iterable_datasets( def _interleave_iterable_datasets( - datasets: list[IterableDataset], - probabilities: Optional[list[float]] = None, + datasets: Sequence[IterableDataset], + probabilities: Optional[Sequence[float]] = None, seed: Optional[int] = None, info: Optional[DatasetInfo] = None, split: Optional[NamedSplit] = None,