# ${generated_comment}
# mypy: allow-untyped-defs
# This base template ("datapipe.pyi.in") is generated from mypy stubgen with minimal editing for code injection
# The output file will be "datapipe.pyi". This is executed as part of torch/CMakeLists.txt
# Note that, for mypy, .pyi file takes precedent over .py file, such that we must define the interface for other
# classes/objects here, even though we are not injecting extra code into them at the moment.

from collections.abc import Iterable, Iterator
from typing import Any, Callable, Literal, Optional, TypeVar, Union

from torch.utils.data import Dataset, default_collate, IterableDataset
from torch.utils.data.datapipes._hook_iterator import _SnapshotState
from torch.utils.data.datapipes._typing import _DataPipeMeta, _IterDataPipeMeta

_T = TypeVar("_T")
_T_co = TypeVar("_T_co", covariant=True)
UNTRACABLE_DATAFRAME_PIPES: Any

class DataChunk(list[_T]):
    items: list[_T]
    def __init__(self, items: Iterable[_T]) -> None: ...
    def as_str(self, indent: str = "") -> str: ...
    def __iter__(self) -> Iterator[_T]: ...
    def raw_iterator(self) -> Iterator[_T]: ...

class MapDataPipe(Dataset[_T_co], metaclass=_DataPipeMeta):
    functions: dict[str, Callable] = ...
    reduce_ex_hook: Callable | None = ...
    getstate_hook: Callable | None = ...
    str_hook: Callable | None = ...
    repr_hook: Callable | None = ...
    def __getattr__(self, attribute_name: Any): ...
    @classmethod
    def register_function(cls, function_name: Any, function: Any) -> None: ...
    @classmethod
    def register_datapipe_as_function(
        cls,
        function_name: Any,
        cls_to_register: Any,
    ): ...
    def __getstate__(self): ...
    def __reduce_ex__(self, *args: Any, **kwargs: Any): ...
    @classmethod
    def set_getstate_hook(cls, hook_fn: Any) -> None: ...
    @classmethod
    def set_reduce_ex_hook(cls, hook_fn: Any) -> None: ...
    ${MapDataPipeMethods}

class IterDataPipe(IterableDataset[_T_co], metaclass=_IterDataPipeMeta):
    functions: dict[str, Callable] = ...
    reduce_ex_hook: Optional[Callable] = ...
    getstate_hook: Optional[Callable] = ...
    str_hook: Optional[Callable] = ...
    repr_hook: Optional[Callable] = ...
    _number_of_samples_yielded: int = ...
    _snapshot_state: _SnapshotState = _SnapshotState.Iterating  # noqa: PYI015
    _fast_forward_iterator: Optional[Iterator] = ...
    def __getattr__(self, attribute_name: Any): ...
    @classmethod
    def register_function(cls, function_name: Any, function: Any) -> None: ...
    @classmethod
    def register_datapipe_as_function(
        cls,
        function_name: Any,
        cls_to_register: Any,
        enable_df_api_tracing: bool = ...,
    ): ...
    def __getstate__(self): ...
    def __reduce_ex__(self, *args: Any, **kwargs: Any): ...
    @classmethod
    def set_getstate_hook(cls, hook_fn: Any) -> None: ...
    @classmethod
    def set_reduce_ex_hook(cls, hook_fn: Any) -> None: ...
    ${IterDataPipeMethods}

class DFIterDataPipe(IterDataPipe):
    def _is_dfpipe(self): ...
    def __iter__(self): ...

class _DataPipeSerializationWrapper:
    def __init__(self, datapipe): ...
    def __getstate__(self): ...
    def __setstate__(self, state): ...
    def __len__(self): ...

class _IterDataPipeSerializationWrapper(_DataPipeSerializationWrapper, IterDataPipe):
    def __iter__(self): ...

class _MapDataPipeSerializationWrapper(_DataPipeSerializationWrapper, MapDataPipe):
    def __getitem__(self, idx): ...