xref: /aosp_15_r20/external/pytorch/torch/_dynamo/polyfills/builtins.py (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1"""
2Python polyfills for builtins
3"""
4
5from __future__ import annotations
6
7import builtins
8from typing import Iterable, TypeVar
9
10from ..decorators import substitute_in_graph
11
12
13__all__ = [
14    "all",
15    "any",
16    "enumerate",
17]
18
19
20_T = TypeVar("_T")
21
22
23@substitute_in_graph(builtins.all, can_constant_fold_through=True)
24def all(iterable: Iterable[object], /) -> bool:
25    for elem in iterable:
26        if not elem:
27            return False
28    return True
29
30
31@substitute_in_graph(builtins.any, can_constant_fold_through=True)
32def any(iterable: Iterable[object], /) -> bool:
33    for elem in iterable:
34        if elem:
35            return True
36    return False
37
38
39@substitute_in_graph(builtins.enumerate, is_embedded_type=True)  # type: ignore[arg-type]
40def enumerate(iterable: Iterable[_T], start: int = 0) -> Iterable[tuple[int, _T]]:
41    if not isinstance(start, int):
42        raise TypeError(
43            f"{type(start).__name__!r} object cannot be interpreted as an integer"
44        )
45
46    for x in iterable:
47        yield start, x
48        start += 1
49