[BE]: Update Typeguard to TypeIs for better type inference (#133814)

Skylion007 · pytorchmergebot · commit 16caa8c1b3a0 · 2024-10-21T17:20:06.000Z
Uses TypeIs instead of TypeGuard for better inference. See https://peps.python.org/pep-0742/ Pull Request resolved: #133814 Approved by: https://github.com/ezyang
diff --git a/.ci/docker/requirements-ci.txt b/.ci/docker/requirements-ci.txt
@@ -257,7 +257,7 @@ tb-nightly==2.13.0a20230426
 #test that import:
 
 # needed by torchgen utils
-typing-extensions
+typing-extensions>=4.10.0
 #Description: type hints for python
 #Pinned versions:
 #test that import:
diff --git a/pyproject.toml b/pyproject.toml
@@ -7,7 +7,7 @@ requires = [
     "ninja",
     "pyyaml",
     "cmake",
-    "typing-extensions",
+    "typing-extensions>=4.10.0",
     "requests",
 ]
 # Use legacy backend to import local packages in setup.py
diff --git a/requirements.txt b/requirements.txt
@@ -11,7 +11,7 @@ requests
 # is required until pytorch build not refactored to work for latest setuptools.
 setuptools<=72.1.0
 types-dataclasses
-typing-extensions>=4.8.0
+typing-extensions>=4.10.0
 sympy==1.13.1 ; python_version >= "3.9"
 filelock
 networkx
diff --git a/setup.py b/setup.py
@@ -1159,7 +1159,7 @@ def main():
         )
     install_requires = [
         "filelock",
-        "typing-extensions>=4.8.0",
+        "typing-extensions>=4.10.0",
         'setuptools ; python_version >= "3.12"',
         'sympy==1.13.1 ; python_version >= "3.9"',
         "networkx",
diff --git a/torch/__init__.py b/torch/__init__.py
@@ -34,7 +34,7 @@
     TypeVar as _TypeVar,
     Union as _Union,
 )
-from typing_extensions import ParamSpec as _ParamSpec, TypeGuard as _TypeGuard
+from typing_extensions import ParamSpec as _ParamSpec, TypeIs as _TypeIs
 
 
 if TYPE_CHECKING:
@@ -1008,7 +1008,7 @@ def typename(obj: _Any, /) -> str:
     return f"{module}.{qualname}"
 
 
-def is_tensor(obj: _Any, /) -> _TypeGuard["torch.Tensor"]:
+def is_tensor(obj: _Any, /) -> _TypeIs["torch.Tensor"]:
     r"""Returns True if `obj` is a PyTorch tensor.
 
     Note that this function is simply doing ``isinstance(obj, Tensor)``.
@@ -1028,7 +1028,7 @@ def is_tensor(obj: _Any, /) -> _TypeGuard["torch.Tensor"]:
     return isinstance(obj, torch.Tensor)
 
 
-def is_storage(obj: _Any, /) -> _TypeGuard[_Union["TypedStorage", "UntypedStorage"]]:
+def is_storage(obj: _Any, /) -> _TypeIs[_Union["TypedStorage", "UntypedStorage"]]:
     r"""Returns True if `obj` is a PyTorch storage object.
 
     Args:
diff --git a/torch/_dynamo/utils.py b/torch/_dynamo/utils.py
@@ -56,7 +56,7 @@
     Union,
     ValuesView,
 )
-from typing_extensions import Literal, TypeGuard
+from typing_extensions import Literal, TypeIs
 
 import torch
 import torch._functorch.config
@@ -569,14 +569,14 @@ def clear(self):
 
 
 @overload
-def istype(obj: object, allowed_types: Type[T]) -> TypeGuard[T]:
+def istype(obj: object, allowed_types: Type[T]) -> TypeIs[T]:
     ...
 
 
 @overload
 def istype(
     obj: object, allowed_types: Tuple[Type[List[T]], Type[Tuple[T, ...]]]
-) -> TypeGuard[T]:
+) -> TypeIs[T]:
     ...
 
 
diff --git a/torch/_inductor/pattern_matcher.py b/torch/_inductor/pattern_matcher.py
@@ -70,7 +70,7 @@
     TypeVar,
     Union,
 )
-from typing_extensions import Self, TypeGuard
+from typing_extensions import Self, TypeIs
 
 import torch
 import torch._guards
@@ -305,10 +305,10 @@ def __bool__(self) -> bool:
 MatchResult = Union[Match, FailedMatch]
 
 
-def is_match(m: MatchResult) -> TypeGuard[Match]:
+def is_match(m: MatchResult) -> TypeIs[Match]:
     """
-    TypeGuards cannot act on `self`. Thus this function exists to let mypy
-    recognize FailedMatch.__bool__ as a TypeGuard.
+    TypeIs cannot act on `self`. Thus this function exists to let mypy
+    recognize FailedMatch.__bool__ as a TypeIs.
     """
     return bool(m)
 
diff --git a/torch/_subclasses/fake_tensor.py b/torch/_subclasses/fake_tensor.py
@@ -32,7 +32,7 @@
     TypeVar,
     Union,
 )
-from typing_extensions import Self, TypeGuard
+from typing_extensions import Self, TypeIs
 from weakref import ReferenceType
 
 import torch
@@ -169,7 +169,7 @@ def get_plain_tensors(subclass: Tensor) -> List[Tensor]:
     return plain_tensors
 
 
-def is_fake(x: object) -> TypeGuard[Tensor]:
+def is_fake(x: object) -> TypeIs[Tensor]:
     if isinstance(x, FakeTensor):
         return True
     if is_traceable_wrapper_subclass(x):
@@ -1213,7 +1213,7 @@ def reset_nt_tensor_id_counter(self) -> None:
     # In this case, it's insufficient to test only one FakeTensor: you need
     # to distinguish between our fake tensor and other fake tensors.  That's
     # what this function does.
-    def is_our_fake(self, t: object) -> TypeGuard[FakeTensor]:
+    def is_our_fake(self, t: object) -> TypeIs[FakeTensor]:
         return isinstance(t, FakeTensor) and t.fake_mode is self
 
     # If we should avoid device init. This changes the behavior of various APIs:
diff --git a/torch/masked/maskedtensor/core.py b/torch/masked/maskedtensor/core.py
@@ -3,7 +3,7 @@
 
 import warnings
 from typing import Any
-from typing_extensions import TypeGuard
+from typing_extensions import TypeIs
 
 import torch
 from torch.overrides import get_default_nowrap_functions
@@ -15,7 +15,7 @@
 ]
 
 
-def is_masked_tensor(obj: Any, /) -> TypeGuard["MaskedTensor"]:
+def is_masked_tensor(obj: Any, /) -> TypeIs["MaskedTensor"]:
     r"""Returns True if the input is a MaskedTensor, else False
 
     Args:
diff --git a/torch/nn/parameter.pyi b/torch/nn/parameter.pyi
@@ -1,5 +1,5 @@
 # mypy: allow-untyped-defs
-from typing_extensions import TypeGuard
+from typing_extensions import TypeIs
 
 from torch import device, dtype, Tensor
 
@@ -8,7 +8,7 @@ class Parameter(Tensor):
 
 def is_lazy(
     param: Tensor,
-) -> TypeGuard[UninitializedParameter | UninitializedBuffer]: ...
+) -> TypeIs[UninitializedParameter | UninitializedBuffer]: ...
 
 class UninitializedParameter(Tensor):
     def __init__(self, data: Tensor = ..., requires_grad: bool = ...) -> None: ...
diff --git a/torch/serialization.py b/torch/serialization.py
@@ -28,7 +28,7 @@
     Type,
     Union,
 )
-from typing_extensions import TypeAlias, TypeGuard  # Python 3.10+
+from typing_extensions import TypeAlias, TypeIs
 
 import torch
 import torch._weights_only_unpickler as _weights_only_unpickler
@@ -620,7 +620,7 @@ def storage_to_tensor_type(storage):
     return getattr(module, storage_type.__name__.replace("Storage", "Tensor"))
 
 
-def _is_path(name_or_buffer) -> TypeGuard[Union[str, os.PathLike]]:
+def _is_path(name_or_buffer) -> TypeIs[Union[str, os.PathLike]]:
     return isinstance(name_or_buffer, (str, os.PathLike))
 
 
diff --git a/torch/utils/_python_dispatch.py b/torch/utils/_python_dispatch.py
@@ -4,7 +4,7 @@
 import warnings
 from dataclasses import dataclass
 from typing import Any, Dict, List, Optional, Set, Union, Protocol, Tuple, Sequence, overload, Deque, Type
-from typing_extensions import TypeGuard
+from typing_extensions import TypeIs
 from collections import deque
 
 import torch
@@ -365,7 +365,7 @@ def to(
 
 
 
-def is_traceable_wrapper_subclass(t: object) -> TypeGuard[TensorWithFlatten]:
+def is_traceable_wrapper_subclass(t: object) -> TypeIs[TensorWithFlatten]:
     """
     Returns whether or not a tensor subclass that implements __torch_dispatch__
     is 'traceable' with torch.compile.
@@ -402,7 +402,7 @@ def is_traceable_wrapper_subclass(t: object) -> TypeGuard[TensorWithFlatten]:
         and hasattr(t, "__tensor_unflatten__")
     )
 
-def is_traceable_wrapper_subclass_type(t: Type) -> TypeGuard[Type[TensorWithFlatten]]:
+def is_traceable_wrapper_subclass_type(t: Type) -> TypeIs[Type[TensorWithFlatten]]:
     """Same as above, but takes a type argument instead of an instance."""
     return (issubclass(t, torch.Tensor) and t != torch.Tensor
             and hasattr(t, "__tensor_flatten__") and hasattr(t, "__tensor_unflatten__"))

Original file line number	Diff line number	Diff line change
`@@ -7,7 +7,7 @@ requires = [`
`7`	`7`	`"ninja",`
`8`	`8`	`"pyyaml",`
`9`	`9`	`"cmake",`
`10`		`- "typing-extensions",`
	`10`	`+ "typing-extensions>=4.10.0",`
`11`	`11`	`"requests",`
`12`	`12`	`]`
`13`	`13`	`# Use legacy backend to import local packages in setup.py`
Original file line number	Diff line number	Diff line change
`@@ -1159,7 +1159,7 @@ def main():`
`1159`	`1159`	`)`
`1160`	`1160`	`install_requires = [`
`1161`	`1161`	`"filelock",`
`1162`		`- "typing-extensions>=4.8.0",`
	`1162`	`+ "typing-extensions>=4.10.0",`
`1163`	`1163`	`'setuptools ; python_version >= "3.12"',`
`1164`	`1164`	`'sympy==1.13.1 ; python_version >= "3.9"',`
`1165`	`1165`	`"networkx",`