OSSForks/pytorch
1
0
Fork 0
mirror of https://github.com/zebrajr/pytorch.git synced 2025-12-06 12:20:52 +01:00
Code Issues Actions 38 Packages Projects Releases Wiki Activity

added type annotation to _NoParamDecoratorContextManager.__new__ (#166414)

Browse Source 
Fixes #166413

Pull Request resolved: https://github.com/pytorch/pytorch/pull/166414
Approved by: https://github.com/Skylion007, https://github.com/malfet
This commit is contained in:
Randolf Scholz 2025-10-28 21:59:17 +00:00 committed by PyTorch MergeBot
parent 84a2715d34
commit fea819ed08
1 changed files with 8 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
torch/utils/_contextlib.py
View File

@ -7,7 +7,8 @@ import inspect
import sys
import warnings
from collections.abc import Callable
from typing import Any, cast, TypeVar
from typing import Any, cast, overload, TypeVar
from typing_extensions import Self
# Used for annotating the decorator usage of _DecoratorContextManager (e.g.,
@ -158,7 +159,12 @@ class _DecoratorContextManager:
class _NoParamDecoratorContextManager(_DecoratorContextManager):
"""Allow a context manager to be used as a decorator without parentheses."""
def __new__(cls, orig_func=None):
@overload
def __new__(cls, orig_func: F) -> F: ... # type: ignore[misc]
@overload
def __new__(cls, orig_func: None = None) -> Self: ...
def __new__(cls, orig_func: F | None = None) -> Self | F: # type: ignore[misc]
if orig_func is None:
return super().__new__(cls)
return cls()(orig_func)