[API Compatibility] Add decorator for paddle.utils.data.BatchSampler (PaddlePaddle#78382)

* add decorator for BatchSampler

* fix error and add first arg type checking

* fix condition

* fix duplicate args

* fix args

* remove duplicate code

* support Iterable parsing and fix tests

Author: algorithm1832

python/paddle/io/dataloader/batch_sampler.py

@@ -16,9 +16,12 @@
 
 import math
 from collections.abc import Iterable, Iterator, Sequence, Sized
+from typing import overload
 
 import numpy as np
 
+from paddle.utils.decorator_utils import batch_sampler_decorator
+
 from .dataset import IterableDataset
 from .sampler import RandomSampler, Sampler, SequenceSampler
 
@@ -107,6 +110,25 @@ class BatchSampler(Sampler[Sequence[int]]):
     shuffle: bool
     drop_last: bool
 
+    @overload
+    def __init__(
+        self,
+        dataset: Sized | None = None,
+        sampler: Sampler | Iterable[int] | None = None,
+        shuffle: bool = False,
+        batch_size: int = 1,
+        drop_last: bool = False,
+    ) -> None: ...
+
+    @overload
+    def __init__(
+        self,
+        sampler: Sampler | Iterable[int] | None = None,
+        batch_size: int = 1,
+        drop_last: bool = False,
+    ) -> None: ...
+
+    @batch_sampler_decorator()
     def __init__(
         self,
         dataset: Sized | None = None,

python/paddle/utils/decorator_utils.py

@@ -17,15 +17,13 @@
 import functools
 import inspect
 import warnings
-from typing import TYPE_CHECKING, Any, Callable, TypeVar, cast
+from collections.abc import Iterable
+from typing import Any, Callable, TypeVar, cast
 
 from typing_extensions import ParamSpec, get_overloads
 
 import paddle
 
-if TYPE_CHECKING:
-    from collections.abc import Iterable
-
 _InputT = ParamSpec("_InputT")
 _RetT = TypeVar("_RetT")
 _SENTINEL = object()
@@ -994,3 +992,39 @@ def wrapper(*args: _InputT.args, **kwargs: _InputT.kwargs) -> _RetT:
         return wrapper
 
     return decorator
+
+
+def batch_sampler_decorator() -> Callable[
+    [Callable[_InputT, _RetT]], Callable[_InputT, _RetT]
+]:
+    """
+    Usage Example:
+    PyTorch: torch.utils.data.BatchSampler(sampler, batch_size, drop_last)
+    Paddle: paddle.utils.data.BatchSampler(dataset, sampler, shuffle, batch_size, drop_last)
+    """
+
+    def decorator(func: Callable[_InputT, _RetT]) -> Callable[_InputT, _RetT]:
+        @functools.wraps(func)
+        def wrapper(*args: _InputT.args, **kwargs: _InputT.kwargs) -> _RetT:
+            # args[0] is self
+            # args[1] is Sampler / Iterable, use torch signature
+            if len(args) >= 2 and isinstance(
+                args[1], (paddle.io.Sampler, Iterable)
+            ):
+                kwargs["sampler"] = args[1]
+                if len(args) >= 3:
+                    kwargs["batch_size"] = args[2]
+                if len(args) == 4:
+                    kwargs["drop_last"] = args[3]
+                if len(args) > 4:
+                    raise TypeError(
+                        "BatchSampler() received too many arguments"
+                    )
+                args = (args[0],)
+
+            return func(*args, **kwargs)
+
+        wrapper.__signature__ = inspect.signature(func)
+        return wrapper
+
+    return decorator

test/legacy_test/test_batch_sampler.py

@@ -225,6 +225,38 @@ def init_batch_sampler(self):
         return bs
 
 
+class TestBatchSamplerTorchPositionalArg(TestBatchSampler):
+    def init_batch_sampler(self):
+        dataset = RandomDataset(1000, 10)
+        sampler = SequenceSampler(dataset)
+        bs = BatchSampler(sampler, self.batch_size, self.drop_last)
+        return bs
+
+
+class TestBatchSamplerTorchPositionalArgWithIterableSampler(TestBatchSampler):
+    def init_batch_sampler(self):
+        sampler = range(1000)
+        bs = BatchSampler(sampler, self.batch_size, self.drop_last)
+        return bs
+
+
+class TestBatchSamplerPositionalArgError(TestBatchSampler):
+    def init_batch_sampler(self):
+        dataset = RandomDataset(1000, 10)
+        sampler = SequenceSampler(dataset)
+        bs = BatchSampler(
+            sampler, self.batch_size, self.drop_last, self.shuffle
+        )
+        return bs
+
+    def test_main(self):
+        try:
+            bs = self.init_batch_sampler()
+            self.assertTrue(False)
+        except TypeError:
+            pass
+
+
 class TestBatchSamplerWithSamplerDropLast(unittest.TestCase):
     def setUp(self):
         self.num_samples = 1000