PaddlePaddle
diff --git a/‎python/paddle/io/dataloader/dataset.py
Copy file name to clipboard
+70-41Lines changed: 70 additions & 41 deletions b/‎python/paddle/io/dataloader/dataset.py
Copy file name to clipboard
+70-41Lines changed: 70 additions & 41 deletions
diff --git a/‎python/paddle/vision/datasets/cifar.py
Copy file name to clipboard
+2-2Lines changed: 2 additions & 2 deletions b/‎python/paddle/vision/datasets/cifar.py
Copy file name to clipboard
+2-2Lines changed: 2 additions & 2 deletions
diff --git a/‎python/paddle/vision/datasets/flowers.py
Copy file name to clipboard
+2-2Lines changed: 2 additions & 2 deletions b/‎python/paddle/vision/datasets/flowers.py
Copy file name to clipboard
+2-2Lines changed: 2 additions & 2 deletions
@@ -12,17 +12,38 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 import bisect
 import math
 import warnings
-from typing import Iterable
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Generator,
+    Generic,
+    Iterable,
+    Iterator,
+    Sequence,
+    Tuple,
+    TypeVar,
+)
+
+from typing_extensions import Never, TypeVarTuple, Unpack
 
 import paddle
 
 from ... import framework
 
+if TYPE_CHECKING:
+    from paddle import Tensor
+
+_T = TypeVar('_T')
+_Ts = TypeVarTuple('_Ts')
+
 
-class Dataset:
+class Dataset(Generic[_T]):
     """
     An abstract class to encapsulate methods and behaviors of datasets.
 
@@ -64,23 +85,23 @@ class Dataset:
             ...     # do something
     """
 
-    def __init__(self):
+    def __init__(self) -> None:
         pass
 
-    def __getitem__(self, idx):
+    def __getitem__(self, idx: int) -> _T:
         raise NotImplementedError(
             "'{}' not implement in class "
             "{}".format('__getitem__', self.__class__.__name__)
         )
 
-    def __len__(self):
+    def __len__(self) -> int:
         raise NotImplementedError(
             "'{}' not implement in class "
             "{}".format('__len__', self.__class__.__name__)
         )
 
 
-class IterableDataset(Dataset):
+class IterableDataset(Dataset[_T]):
     """
     An abstract class to encapsulate methods and behaviors of iterable datasets.
 
@@ -241,29 +262,29 @@ class IterableDataset(Dataset):
 
     """
 
-    def __init__(self):
+    def __init__(self) -> None:
         pass
 
-    def __iter__(self):
+    def __iter__(self) -> Iterator[_T]:
         raise NotImplementedError(
             "'{}' not implement in class "
             "{}".format('__iter__', self.__class__.__name__)
         )
 
-    def __getitem__(self, idx):
+    def __getitem__(self, idx: int) -> Never:
         raise RuntimeError(
             "'{}' should not be called for IterableDataset"
             "{}".format('__getitem__', self.__class__.__name__)
         )
 
-    def __len__(self):
+    def __len__(self) -> Never:
         raise RuntimeError(
             "'{}' should not be called for IterableDataset"
             "{}".format('__len__', self.__class__.__name__)
         )
 
 
-class TensorDataset(Dataset):
+class TensorDataset(Dataset["Tensor"]):
     """
     Dataset defined by a list of tensors.
 
@@ -298,7 +319,9 @@ class TensorDataset(Dataset):
             ...     # do something
     """
 
-    def __init__(self, tensors):
+    tensors: Sequence[Tensor]
+
+    def __init__(self, tensors: Sequence[Tensor]) -> None:
         if not framework.in_dynamic_mode():
             raise RuntimeError(
                 "TensorDataset con only be used in imperative mode"
@@ -308,10 +331,10 @@ def __init__(self, tensors):
         ), "tensors not have same shape of the 1st dimension"
         self.tensors = tensors
 
-    def __getitem__(self, index):
+    def __getitem__(self, index: int) -> tuple[Tensor, ...]:
         return tuple(tensor[index] for tensor in self.tensors)
 
-    def __len__(self):
+    def __len__(self) -> int:
         return self.tensors[0].shape[0]
 
 
@@ -323,7 +346,7 @@ def to_list(value):
     return [value]
 
 
-class ComposeDataset(Dataset):
+class ComposeDataset(Dataset[Tuple[Unpack[_Ts]]]):
     """
     A Dataset which composes fields of multiple datasets.
 
@@ -344,9 +367,8 @@ class ComposeDataset(Dataset):
             >>> import paddle
             >>> from paddle.io import Dataset, ComposeDataset
 
-
             >>> # define a random dataset
-            >>> class RandomDataset(Dataset):
+            >>> class RandomDataset(Dataset):  # type: ignore[type-arg]
             ...     def __init__(self, num_samples):
             ...         self.num_samples = num_samples
             ...
@@ -358,13 +380,15 @@ class ComposeDataset(Dataset):
             ...     def __len__(self):
             ...         return self.num_samples
             ...
-            >>> dataset = ComposeDataset([RandomDataset(10), RandomDataset(10)])
+            >>> dataset = ComposeDataset([RandomDataset(10), RandomDataset(10)])  # type: ignore[var-annotated]
             >>> for i in range(len(dataset)):
             ...     image1, label1, image2, label2 = dataset[i]
             ...     # do something
     """
 
-    def __init__(self, datasets):
+    datasets: list[Dataset[Any]]
+
+    def __init__(self, datasets: list[Dataset[Any]]) -> None:
         self.datasets = list(datasets)
         assert len(self.datasets) > 0, "input datasets should not be empty"
         for i, dataset in enumerate(self.datasets):
@@ -379,17 +403,17 @@ def __init__(self, datasets):
                     self.datasets[i - 1]
                 ), "lengths of datasets should be same"
 
-    def __len__(self):
+    def __len__(self) -> int:
         return len(self.datasets[0])
 
-    def __getitem__(self, idx):
+    def __getitem__(self, idx) -> tuple[Unpack[_Ts]]:
         sample = []
         for dataset in self.datasets:
             sample.extend(to_list(dataset[idx]))
         return tuple(sample)
 
 
-class ChainDataset(IterableDataset):
+class ChainDataset(IterableDataset[Any]):
     """
     A Dataset which chains multiple iterable-style datasets.
 
@@ -429,20 +453,20 @@ class ChainDataset(IterableDataset):
 
     """
 
-    def __init__(self, datasets):
+    def __init__(self, datasets: list[IterableDataset[Any]]):
         self.datasets = list(datasets)
         assert len(self.datasets) > 0, "input datasets should not be empty"
         for i, dataset in enumerate(self.datasets):
             assert isinstance(
                 dataset, IterableDataset
             ), "ChainDataset only support paddle.io.IterableDataset"
 
-    def __iter__(self):
+    def __iter__(self) -> Iterator[Any]:
         for dataset in self.datasets:
             yield from dataset
 
 
-class Subset(Dataset):
+class Subset(Dataset[_T]):
     """
     Subset of a dataset at specified indices.
 
@@ -471,18 +495,22 @@ class Subset(Dataset):
             [2, 2]
     """
 
-    def __init__(self, dataset, indices):
+    def __init__(self, dataset: Dataset[_T], indices: Sequence[int]) -> None:
         self.dataset = dataset
         self.indices = indices
 
-    def __getitem__(self, idx):
+    def __getitem__(self, idx: int) -> _T:
         return self.dataset[self.indices[idx]]
 
-    def __len__(self):
+    def __len__(self) -> int:
         return len(self.indices)
 
 
-def random_split(dataset, lengths, generator=None):
+def random_split(
+    dataset: Dataset[_T],
+    lengths: Sequence[int],
+    generator: Any | None = None,
+) -> list[Subset[_T]]:
     """
     Randomly split a dataset into non-overlapping new datasets of given lengths.
     Optionally fix the generator for reproducible results, e.g.:
@@ -502,19 +530,19 @@ def random_split(dataset, lengths, generator=None):
             >>> import paddle
 
             >>> paddle.seed(2023)
-            >>> a_list = paddle.io.random_split(range(10), [3, 7])
+            >>> a_list = paddle.io.random_split(range(10), [3, 7])  # type: ignore[arg-type, var-annotated]
             >>> print(len(a_list))
             2
 
             >>> # output of the first subset
-            >>> for idx, v in enumerate(a_list[0]):
+            >>> for idx, v in enumerate(a_list[0]):  # type: ignore[arg-type, var-annotated]
             ...     print(idx, v) # doctest: +SKIP("The output depends on the environment.")
             0 7
             1 6
             2 5
 
             >>> # output of the second subset
-            >>> for idx, v in enumerate(a_list[1]):
+            >>> for idx, v in enumerate(a_list[1]):  # type: ignore[arg-type, var-annotated]
             ...     print(idx, v) # doctest: +SKIP("The output depends on the environment.")
             0 1
             1 9
@@ -560,7 +588,9 @@ def random_split(dataset, lengths, generator=None):
     ]
 
 
-def _accumulate(iterable, fn=lambda x, y: x + y):
+def _accumulate(
+    iterable: Iterable[_T], fn: Callable[[_T, _T], _T] = lambda x, y: x + y
+) -> Generator[_T, None, None]:
     """
     Return running totals
 
@@ -595,7 +625,7 @@ def _accumulate(iterable, fn=lambda x, y: x + y):
         yield total
 
 
-class ConcatDataset(Dataset):
+class ConcatDataset(Dataset[_T]):
     """
     Dataset as a concatenation of multiple datasets.
 
@@ -615,9 +645,8 @@ class ConcatDataset(Dataset):
             >>> import paddle
             >>> from paddle.io import Dataset, ConcatDataset
 
-
             >>> # define a random dataset
-            >>> class RandomDataset(Dataset):
+            >>> class RandomDataset(Dataset):  # type: ignore[type-arg]
             ...     def __init__(self, num_samples):
             ...         self.num_samples = num_samples
             ...
@@ -629,22 +658,22 @@ class ConcatDataset(Dataset):
             ...     def __len__(self):
             ...         return self.num_samples
             ...
-            >>> dataset = ConcatDataset([RandomDataset(10), RandomDataset(10)])
+            >>> dataset = ConcatDataset([RandomDataset(10), RandomDataset(10)])  # type: ignore[var-annotated]
             >>> for i in range(len(dataset)):
             ...     image, label = dataset[i]
             ...     # do something
     """
 
     @staticmethod
-    def cumsum(sequence):
+    def cumsum(sequence: Sequence[Any]) -> list[int]:
         r, s = [], 0
         for e in sequence:
             l = len(e)
             r.append(l + s)
             s += l
         return r
 
-    def __init__(self, datasets: Iterable[Dataset]):
+    def __init__(self, datasets: Iterable[Dataset[Any]]) -> None:
         self.datasets = list(datasets)
         assert (
             len(self.datasets) > 0
@@ -655,10 +684,10 @@ def __init__(self, datasets: Iterable[Dataset]):
             ), "ConcatDataset does not support IterableDataset"
         self.cumulative_sizes = self.cumsum(self.datasets)
 
-    def __len__(self):
+    def __len__(self) -> int:
         return self.cumulative_sizes[-1]
 
-    def __getitem__(self, idx):
+    def __getitem__(self, idx: int) -> _T:
         if idx < 0:
             if -idx > len(self):
                 raise ValueError(
 
@@ -13,7 +13,7 @@
 # limitations under the License.
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, Literal
+from typing import TYPE_CHECKING, Any, Literal, Tuple
 
 if TYPE_CHECKING:
     import numpy.typing as npt
@@ -52,7 +52,7 @@
 }
 
 
-class Cifar10(Dataset):
+class Cifar10(Dataset[Tuple["_ImageDataType", "npt.NDArray[Any]"]]):
     """
     Implementation of `Cifar-10 <https://www.cs.toronto.edu/~kriz/cifar.html>`_
     dataset, which has 10 categories.
 
@@ -14,7 +14,7 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, Literal
+from typing import TYPE_CHECKING, Any, Literal, Tuple
 
 if TYPE_CHECKING:
     import numpy.typing as npt
@@ -51,7 +51,7 @@
 MODE_FLAG_MAP = {'train': 'tstid', 'test': 'trnid', 'valid': 'valid'}
 
 
-class Flowers(Dataset):
+class Flowers(Dataset[Tuple["_ImageDataType", "npt.NDArray[np.int64]"]]):
     """
     Implementation of `Flowers102 <https://www.robots.ox.ac.uk/~vgg/data/flowers/>`_
     dataset.