mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 12:54:11 +08:00
3cda34ebde
This is the result of applying the ruff `UP035` check. `Callable` is imported from `collections.abc` instead of `typing`. `TypeAlias` is also imported from `typing`. This PR is the follow-up of #163947. Pull Request resolved: https://github.com/pytorch/pytorch/pull/164054 Approved by: https://github.com/ezyang, https://github.com/Skylion007
30 lines
652 B
Python
30 lines
652 B
Python
from collections.abc import Callable
|
|
from typing import Generic, Optional, TypeVar
|
|
|
|
|
|
R = TypeVar("R")
|
|
|
|
|
|
class Thunk(Generic[R]):
|
|
"""
|
|
A simple lazy evaluation implementation that lets you delay
|
|
execution of a function. It properly handles releasing the
|
|
function once it is forced.
|
|
"""
|
|
|
|
f: Optional[Callable[[], R]]
|
|
r: Optional[R]
|
|
|
|
__slots__ = ["f", "r"]
|
|
|
|
def __init__(self, f: Callable[[], R]):
|
|
self.f = f
|
|
self.r = None
|
|
|
|
def force(self) -> R:
|
|
if self.f is None:
|
|
return self.r # type: ignore[return-value]
|
|
self.r = self.f()
|
|
self.f = None
|
|
return self.r
|