from typing import Any, Iterator, List, TypeVar

from .pytorch import IterableDataset
from .utils import PipelineStage

T = TypeVar('T')

def add_length_method(obj: Any) -> Any: ...

class DataPipeline(IterableDataset, PipelineStage):
    pipeline: List[Any]
    length: int
    repetitions: int
    nsamples: int
    size: int

    def __init__(self, *args: Any, **kwargs: Any) -> None: ...
    def close(self) -> None: ...
    def invoke(self, f: Any, *args: Any, **kwargs: Any) -> Any: ...
    def iterator1(self) -> Iterator[Any]: ...
    def iterator(self) -> Iterator[Any]: ...
    def __iter__(self) -> Iterator[Any]: ...
    def stage(self, i: int) -> Any: ...
    def append(self, f: Any) -> None: ...
    def compose(self, *args: Any) -> 'DataPipeline': ...
    def with_length(self, n: int, silent: bool = False) -> 'DataPipeline': ...
    def with_epoch(self, nsamples: int = -1, nbatches: int = -1) -> 'DataPipeline': ...
    def repeat(self, nepochs: int = -1, nbatches: int = -1) -> 'DataPipeline': ...