Support typing.assert_type.

* Aliases typing.assert_type to our assert_type pseudo-builtin.
* Removes single-argument assert_type form. This functionality didn't make it
  into the typing library variant. I don't think it's used much anyway,
  given that pytype_extensions.assert_type doesn't support it at runtime, and
  no one's ever noticed.

PiperOrigin-RevId: 577000283

Author: rchen152

pytype/errors.py

@@ -1250,16 +1250,10 @@ def reveal_type(self, stack, node, var):
     self.error(stack, self._var_to_printed_type(var, node))
 
   @_error_name("assert-type")
-  def assert_type(self, stack, node, var, typ=None):
+  def assert_type(self, stack, node, var, typ):
     """Check that a variable type matches its expected value."""
     actual = self._var_to_printed_type(var, node)
 
-    # assert_type(x) checks that x is not Any
-    if typ is None:
-      if actual in ("Any", "typing.Any"):
-        self.error(stack, f"Asserted type was {actual}")
-      return
-
     try:
       expected = abstract_utils.get_atomic_python_constant(typ, str)
     except abstract_utils.ConversionError:

pytype/overlays/CMakeLists.txt

@@ -320,6 +320,7 @@ py_library(
     .named_tuple
     .overlay
     .overlay_utils
+    .special_builtins
     .typed_dict
     pytype.utils
     pytype.abstract.abstract

pytype/overlays/special_builtins.py

@@ -553,10 +553,7 @@ class AssertType(BuiltinFunction):
   _NAME = "assert_type"
 
   def call(self, node, func, args, alias_map=None):
-    if len(args.posargs) == 1:
-      a, = args.posargs
-      t = None
-    elif len(args.posargs) == 2:
+    if len(args.posargs) == 2:
       a, t = args.posargs
     else:
       raise function.WrongArgCount(self._SIGNATURE, args, self.ctx)

pytype/overlays/typing_overlay.py

@@ -16,6 +16,7 @@
 from pytype.overlays import named_tuple
 from pytype.overlays import overlay
 from pytype.overlays import overlay_utils
+from pytype.overlays import special_builtins
 from pytype.overlays import typed_dict
 from pytype.pytd import pep484
 from pytype.pytd import pytd
@@ -627,7 +628,6 @@ def build_re_member(ctx, module):
     "TypeVarTuple": (3, 11),
     "Unpack": (3, 11),
     "assert_never": (3, 11),
-    "assert_type": (3, 11),
     "reveal_type": (3, 11),
 }
 
@@ -658,6 +658,9 @@ def build_re_member(ctx, module):
     "TypedDict": (overlay.drop_module(typed_dict.TypedDictBuilder), (3, 8)),
     "Union": (overlay.drop_module(Union), None),
     "TYPE_CHECKING": (overlay.drop_module(build_typechecking), None),
+    "assert_type": (
+        overlay.add_name("assert_type", special_builtins.AssertType.make_alias),
+        (3, 11)),
     "cast": (overlay.add_name("cast", Cast.make), None),
     "clear_overloads": (_builder_from_name("clear_overloads"), (3, 11)),
     "dataclass_transform": (overlay.add_name(

pytype/tests/test_errors2.py

@@ -345,31 +345,27 @@ def test_assert_type(self):
       class A: pass
       def f(x: int, y: str, z):
         assert_type(x, int)
-        assert_type(y, int)  # assert-type[e1]
-        assert_type(z)  # assert-type[e2]
+        assert_type(y, int)  # assert-type[e]
         if __random__:
           x = A()
         assert_type(x, Union[A, int])
     """)
     self.assertErrorSequences(errors, {
-        "e1": ["Expected", "int", "Actual", "str"],
-        "e2": ["type was Any"]
+        "e": ["Expected", "int", "Actual", "str"],
     })
 
   def test_assert_type_str(self):
     _, errors = self.InferWithErrors("""
       class A: pass
       def f(x: int, y: str, z):
         assert_type(x, 'int')
-        assert_type(y, 'int')  # assert-type[e1]
-        assert_type(z)  # assert-type[e2]
+        assert_type(y, 'int')  # assert-type[e]
         if __random__:
           x = A()
         assert_type(x, 'Union[A, int]')
     """)
     self.assertErrorSequences(errors, {
-        "e1": ["Expected", "int", "Actual", "str"],
-        "e2": ["type was Any"]
+        "e": ["Expected", "int", "Actual", "str"],
     })
 
   def test_assert_type_import(self):
@@ -383,15 +379,13 @@ def assert_type(*args): ...
         class A: pass
         def f(x: int, y: str, z):
           assert_type(x, int)
-          assert_type(y, int)  # assert-type[e1]
-          assert_type(z)  # assert-type[e2]
+          assert_type(y, int)  # assert-type[e]
           if __random__:
             x = A()
           assert_type(x, Union[A, int])
       """, pythonpath=[d.path])
       self.assertErrorSequences(errors, {
-          "e1": ["Expected", "int", "Actual", "str"],
-          "e2": ["type was Any"]
+          "e": ["Expected", "int", "Actual", "str"],
       })
 
   def test_combine_containers(self):

pytype/tests/test_pattern_matching.py

@@ -718,8 +718,8 @@ def f(x: A | B, y: A | B):
             assert_type(c, int)
             assert_type(d, str)
           case _:
-            assert_type(x)
-            assert_type(y)
+            assert_type(x, B)
+            assert_type(y, A)
     """)
 
   def test_builtin(self):