From b684d0342334cb3110621ad37cbfea9a0e137fbb Mon Sep 17 00:00:00 2001
From: rechen <rechen@google.com>
Date: Thu, 13 May 2021 16:16:59 -0700
Subject: [PATCH] Allow typing.TypeGuard to be used in type stubs.

For https://github.com/google/pytype/issues/916.

PiperOrigin-RevId: 373679585
---
 pytype/pyi/definitions.py | 14 ++++++++++----
 pytype/pyi/parser_test.py | 29 +++++++++++++++++++++++++++++
 2 files changed, 39 insertions(+), 4 deletions(-)

diff --git a/pytype/pyi/definitions.py b/pytype/pyi/definitions.py
index 51679ffa5..a52b7e01d 100644
--- a/pytype/pyi/definitions.py
+++ b/pytype/pyi/definitions.py
@@ -25,10 +25,12 @@
 _TYPING_SETS = ("typing.Intersection", "typing.Optional", "typing.Union")
 
 # Aliases for some typing.X types
-_TUPLE_TYPES = ("tuple", "builtins.tuple", "typing.Tuple")
+_ANNOTATED_TYPES = ("typing.Annotated", "typing_extensions.Annotated")
 _CALLABLE_TYPES = ("typing.Callable", "collections.abc.Callable")
+_CONCATENATE_TYPES = ("typing.Concatenate", "typing_extensions.Concatenate")
 _LITERAL_TYPES = ("typing.Literal", "typing_extensions.Literal")
-_ANNOTATED_TYPES = ("typing.Annotated", "typing_extensions.Annotated")
+_TUPLE_TYPES = ("tuple", "builtins.tuple", "typing.Tuple")
+_TYPEGUARD_TYPES = ("typing.TypeGuard", "typing_extensions.TypeGuard")
 
 
 def _split_definitions(defs: List[Any]):
@@ -432,6 +434,10 @@ def _parameterized_type(self, base_type, parameters):
       return types.pytd_literal(parameters)
     elif self._matches_named_type(base_type, _ANNOTATED_TYPES):
       return types.pytd_annotated(parameters)
+    elif self._matches_named_type(base_type, _TYPEGUARD_TYPES):
+      # We do not yet support PEP 647, User-Defined Type Guards. To avoid
+      # blocking typeshed, convert type guards to plain bools.
+      return pytd.NamedType("bool")
     elif any(isinstance(p, types.Constant) for p in parameters):
       parameters = ", ".join(
           p.repr_str() if isinstance(p, types.Constant) else "_"
@@ -458,8 +464,8 @@ def _parameterized_type(self, base_type, parameters):
           # To avoid blocking typeshed from adopting this PEP, we convert new
           # features to Any.
           if p in self.param_specs or (
-              isinstance(p, pytd.GenericType) and self._matches_full_name(
-                  p, ("typing.Concatenate", "typing_extensions.Concatenate"))):
+              isinstance(p, pytd.GenericType) and
+              self._matches_full_name(p, _CONCATENATE_TYPES)):
             callable_parameters.append(pytd.AnythingType())
           else:
             callable_parameters.append(p)
diff --git a/pytype/pyi/parser_test.py b/pytype/pyi/parser_test.py
index d2058fc6d..c9193f04e 100644
--- a/pytype/pyi/parser_test.py
+++ b/pytype/pyi/parser_test.py
@@ -2925,5 +2925,34 @@ def f(x: Callable[..., R]) -> Callable[..., R]: ...
     """)
 
 
+class TypeGuardTest(_ParserTestBase):
+
+  def test_typing_extensions(self):
+    self.check("""
+      from typing import List
+
+      from typing_extensions import TypeGuard
+
+      def f(x: List[object]) -> TypeGuard[List[str]]: ...
+    """, """
+      from typing import List
+
+      from typing_extensions import TypeGuard
+
+      def f(x: List[object]) -> bool: ...
+  """)
+
+  def test_typing(self):
+    self.check("""
+      from typing import List, TypeGuard
+
+      def f(x: List[object]) -> TypeGuard[List[str]]: ...
+    """, """
+      from typing import List
+
+      def f(x: List[object]) -> bool: ...
+    """)
+
+
 if __name__ == "__main__":
   unittest.main()