Add support for assert_type (#12584)

JelleZijlstra · web-flow · commit 875774769922 · 2022-04-14T20:53:33.000-07:00
See python/cpython#30843. The implementation mostly follows that of cast(). It relies on `mypy.sametypes.is_same_type()`.
diff --git a/mypy/checkexpr.py b/mypy/checkexpr.py
@@ -23,7 +23,7 @@
     get_proper_types, flatten_nested_unions, LITERAL_TYPE_NAMES,
 )
 from mypy.nodes import (
-    NameExpr, RefExpr, Var, FuncDef, OverloadedFuncDef, TypeInfo, CallExpr,
+    AssertTypeExpr, NameExpr, RefExpr, Var, FuncDef, OverloadedFuncDef, TypeInfo, CallExpr,
     MemberExpr, IntExpr, StrExpr, BytesExpr, UnicodeExpr, FloatExpr,
     OpExpr, UnaryExpr, IndexExpr, CastExpr, RevealExpr, TypeApplication, ListExpr,
     TupleExpr, DictExpr, LambdaExpr, SuperExpr, SliceExpr, Context, Expression,
@@ -3144,6 +3144,14 @@ def visit_cast_expr(self, expr: CastExpr) -> Type:
                                context=expr)
         return target_type
 
+    def visit_assert_type_expr(self, expr: AssertTypeExpr) -> Type:
+        source_type = self.accept(expr.expr, type_context=AnyType(TypeOfAny.special_form),
+                                  allow_none_return=True, always_allow_any=True)
+        target_type = expr.type
+        if not is_same_type(source_type, target_type):
+            self.msg.assert_type_fail(source_type, target_type, expr)
+        return source_type
+
     def visit_reveal_expr(self, expr: RevealExpr) -> Type:
         """Type check a reveal_type expression."""
         if expr.kind == REVEAL_TYPE:
diff --git a/mypy/errorcodes.py b/mypy/errorcodes.py
@@ -113,6 +113,9 @@ def __str__(self) -> str:
 REDUNDANT_CAST: Final = ErrorCode(
     "redundant-cast", "Check that cast changes type of expression", "General"
 )
+ASSERT_TYPE: Final = ErrorCode(
+    "assert-type", "Check that assert_type() call succeeds", "General"
+)
 COMPARISON_OVERLAP: Final = ErrorCode(
     "comparison-overlap", "Check that types in comparisons and 'in' expressions overlap", "General"
 )
diff --git a/mypy/literals.py b/mypy/literals.py
@@ -8,7 +8,8 @@
     ConditionalExpr, EllipsisExpr, YieldFromExpr, YieldExpr, RevealExpr, SuperExpr,
     TypeApplication, LambdaExpr, ListComprehension, SetComprehension, DictionaryComprehension,
     GeneratorExpr, BackquoteExpr, TypeVarExpr, TypeAliasExpr, NamedTupleExpr, EnumCallExpr,
-    TypedDictExpr, NewTypeExpr, PromoteExpr, AwaitExpr, TempNode, AssignmentExpr, ParamSpecExpr
+    TypedDictExpr, NewTypeExpr, PromoteExpr, AwaitExpr, TempNode, AssignmentExpr, ParamSpecExpr,
+    AssertTypeExpr,
 )
 from mypy.visitor import ExpressionVisitor
 
@@ -175,6 +176,9 @@ def visit_slice_expr(self, e: SliceExpr) -> None:
     def visit_cast_expr(self, e: CastExpr) -> None:
         return None
 
+    def visit_assert_type_expr(self, e: AssertTypeExpr) -> None:
+        return None
+
     def visit_conditional_expr(self, e: ConditionalExpr) -> None:
         return None
 
diff --git a/mypy/messages.py b/mypy/messages.py
@@ -1213,6 +1213,11 @@ def redundant_cast(self, typ: Type, context: Context) -> None:
         self.fail('Redundant cast to {}'.format(format_type(typ)), context,
                   code=codes.REDUNDANT_CAST)
 
+    def assert_type_fail(self, source_type: Type, target_type: Type, context: Context) -> None:
+        self.fail(f"Expression is of type {format_type(source_type)}, "
+                  f"not {format_type(target_type)}", context,
+                  code=codes.ASSERT_TYPE)
+
     def unimported_type_becomes_any(self, prefix: str, typ: Type, ctx: Context) -> None:
         self.fail("{} becomes {} due to an unfollowed import".format(prefix, format_type(typ)),
                   ctx, code=codes.NO_ANY_UNIMPORTED)
diff --git a/mypy/mixedtraverser.py b/mypy/mixedtraverser.py
@@ -1,7 +1,7 @@
 from typing import Optional
 
 from mypy.nodes import (
-    Var, FuncItem, ClassDef, AssignmentStmt, ForStmt, WithStmt,
+    AssertTypeExpr, Var, FuncItem, ClassDef, AssignmentStmt, ForStmt, WithStmt,
     CastExpr, TypeApplication, TypeAliasExpr, TypeVarExpr, TypedDictExpr, NamedTupleExpr,
     PromoteExpr, NewTypeExpr
 )
@@ -79,6 +79,10 @@ def visit_cast_expr(self, o: CastExpr) -> None:
         super().visit_cast_expr(o)
         o.type.accept(self)
 
+    def visit_assert_type_expr(self, o: AssertTypeExpr) -> None:
+        super().visit_assert_type_expr(o)
+        o.type.accept(self)
+
     def visit_type_application(self, o: TypeApplication) -> None:
         super().visit_type_application(o)
         for t in o.types:
diff --git a/mypy/nodes.py b/mypy/nodes.py
@@ -1945,6 +1945,22 @@ def accept(self, visitor: ExpressionVisitor[T]) -> T:
         return visitor.visit_cast_expr(self)
 
 
+class AssertTypeExpr(Expression):
+    """Represents a typing.assert_type(expr, type) call."""
+    __slots__ = ('expr', 'type')
+
+    expr: Expression
+    type: "mypy.types.Type"
+
+    def __init__(self, expr: Expression, typ: 'mypy.types.Type') -> None:
+        super().__init__()
+        self.expr = expr
+        self.type = typ
+
+    def accept(self, visitor: ExpressionVisitor[T]) -> T:
+        return visitor.visit_assert_type_expr(self)
+
+
 class RevealExpr(Expression):
     """Reveal type expression reveal_type(expr) or reveal_locals() expression."""
 
diff --git a/mypy/semanal.py b/mypy/semanal.py
@@ -56,7 +56,7 @@
 from typing_extensions import Final, TypeAlias as _TypeAlias
 
 from mypy.nodes import (
-    MypyFile, TypeInfo, Node, AssignmentStmt, FuncDef, OverloadedFuncDef,
+    AssertTypeExpr, MypyFile, TypeInfo, Node, AssignmentStmt, FuncDef, OverloadedFuncDef,
     ClassDef, Var, GDEF, FuncItem, Import, Expression, Lvalue,
     ImportFrom, ImportAll, Block, LDEF, NameExpr, MemberExpr,
     IndexExpr, TupleExpr, ListExpr, ExpressionStmt, ReturnStmt,
@@ -99,7 +99,7 @@
     TypeTranslator, TypeOfAny, TypeType, NoneType, PlaceholderType, TPDICT_NAMES, ProperType,
     get_proper_type, get_proper_types, TypeAliasType, TypeVarLikeType, Parameters, ParamSpecType,
     PROTOCOL_NAMES, TYPE_ALIAS_NAMES, FINAL_TYPE_NAMES, FINAL_DECORATOR_NAMES, REVEAL_TYPE_NAMES,
-    is_named_instance,
+    ASSERT_TYPE_NAMES, is_named_instance,
 )
 from mypy.typeops import function_type, get_type_vars
 from mypy.type_visitor import TypeQuery
@@ -3897,6 +3897,19 @@ def visit_call_expr(self, expr: CallExpr) -> None:
             expr.analyzed.line = expr.line
             expr.analyzed.column = expr.column
             expr.analyzed.accept(self)
+        elif refers_to_fullname(expr.callee, ASSERT_TYPE_NAMES):
+            if not self.check_fixed_args(expr, 2, 'assert_type'):
+                return
+            # Translate second argument to an unanalyzed type.
+            try:
+                target = self.expr_to_unanalyzed_type(expr.args[1])
+            except TypeTranslationError:
+                self.fail('assert_type() type is not a type', expr)
+                return
+            expr.analyzed = AssertTypeExpr(expr.args[0], target)
+            expr.analyzed.line = expr.line
+            expr.analyzed.column = expr.column
+            expr.analyzed.accept(self)
         elif refers_to_fullname(expr.callee, REVEAL_TYPE_NAMES):
             if not self.check_fixed_args(expr, 1, 'reveal_type'):
                 return
@@ -4200,6 +4213,12 @@ def visit_cast_expr(self, expr: CastExpr) -> None:
         if analyzed is not None:
             expr.type = analyzed
 
+    def visit_assert_type_expr(self, expr: AssertTypeExpr) -> None:
+        expr.expr.accept(self)
+        analyzed = self.anal_type(expr.type)
+        if analyzed is not None:
+            expr.type = analyzed
+
     def visit_reveal_expr(self, expr: RevealExpr) -> None:
         if expr.kind == REVEAL_TYPE:
             if expr.expr is not None:
diff --git a/mypy/server/astmerge.py b/mypy/server/astmerge.py
@@ -51,7 +51,7 @@
     MypyFile, SymbolTable, Block, AssignmentStmt, NameExpr, MemberExpr, RefExpr, TypeInfo,
     FuncDef, ClassDef, NamedTupleExpr, SymbolNode, Var, Statement, SuperExpr, NewTypeExpr,
     OverloadedFuncDef, LambdaExpr, TypedDictExpr, EnumCallExpr, FuncBase, TypeAliasExpr, CallExpr,
-    CastExpr, TypeAlias,
+    CastExpr, TypeAlias, AssertTypeExpr,
     MDEF
 )
 from mypy.traverser import TraverserVisitor
@@ -226,6 +226,10 @@ def visit_cast_expr(self, node: CastExpr) -> None:
         super().visit_cast_expr(node)
         self.fixup_type(node.type)
 
+    def visit_assert_type_expr(self, node: AssertTypeExpr) -> None:
+        super().visit_assert_type_expr(node)
+        self.fixup_type(node.type)
+
     def visit_super_expr(self, node: SuperExpr) -> None:
         super().visit_super_expr(node)
         if node.info is not None:
diff --git a/mypy/server/deps.py b/mypy/server/deps.py
@@ -89,7 +89,8 @@ class 'mod.Cls'. This can also refer to an attribute inherited from a
     ComparisonExpr, GeneratorExpr, DictionaryComprehension, StarExpr, PrintStmt, ForStmt, WithStmt,
     TupleExpr, OperatorAssignmentStmt, DelStmt, YieldFromExpr, Decorator, Block,
     TypeInfo, FuncBase, OverloadedFuncDef, RefExpr, SuperExpr, Var, NamedTupleExpr, TypedDictExpr,
-    LDEF, MDEF, GDEF, TypeAliasExpr, NewTypeExpr, ImportAll, EnumCallExpr, AwaitExpr
+    LDEF, MDEF, GDEF, TypeAliasExpr, NewTypeExpr, ImportAll, EnumCallExpr, AwaitExpr,
+    AssertTypeExpr,
 )
 from mypy.operators import (
     op_methods, reverse_op_methods, ops_with_inplace_method, unary_op_methods
@@ -686,6 +687,10 @@ def visit_cast_expr(self, e: CastExpr) -> None:
         super().visit_cast_expr(e)
         self.add_type_dependencies(e.type)
 
+    def visit_assert_type_expr(self, e: AssertTypeExpr) -> None:
+        super().visit_assert_type_expr(e)
+        self.add_type_dependencies(e.type)
+
     def visit_type_application(self, e: TypeApplication) -> None:
         super().visit_type_application(e)
         for typ in e.types:
diff --git a/mypy/server/subexpr.py b/mypy/server/subexpr.py
@@ -7,7 +7,7 @@
     SliceExpr, CastExpr, RevealExpr, UnaryExpr, ListExpr, TupleExpr, DictExpr, SetExpr,
     IndexExpr, GeneratorExpr, ListComprehension, SetComprehension, DictionaryComprehension,
     ConditionalExpr, TypeApplication, LambdaExpr, StarExpr, BackquoteExpr, AwaitExpr,
-    AssignmentExpr,
+    AssignmentExpr, AssertTypeExpr,
 )
 from mypy.traverser import TraverserVisitor
 
@@ -99,6 +99,10 @@ def visit_cast_expr(self, e: CastExpr) -> None:
         self.add(e)
         super().visit_cast_expr(e)
 
+    def visit_assert_type_expr(self, e: AssertTypeExpr) -> None:
+        self.add(e)
+        super().visit_assert_type_expr(e)
+
     def visit_reveal_expr(self, e: RevealExpr) -> None:
         self.add(e)
         super().visit_reveal_expr(e)
diff --git a/mypy/strconv.py b/mypy/strconv.py
@@ -431,6 +431,9 @@ def visit_comparison_expr(self, o: 'mypy.nodes.ComparisonExpr') -> str:
     def visit_cast_expr(self, o: 'mypy.nodes.CastExpr') -> str:
         return self.dump([o.expr, o.type], o)
 
+    def visit_assert_type_expr(self, o: 'mypy.nodes.AssertTypeExpr') -> str:
+        return self.dump([o.expr, o.type], o)
+
     def visit_reveal_expr(self, o: 'mypy.nodes.RevealExpr') -> str:
         if o.kind == mypy.nodes.REVEAL_TYPE:
             return self.dump([o.expr], o)
diff --git a/mypy/traverser.py b/mypy/traverser.py
@@ -9,7 +9,7 @@
 )
 from mypy.visitor import NodeVisitor
 from mypy.nodes import (
-    Block, MypyFile, FuncBase, FuncItem, CallExpr, ClassDef, Decorator, FuncDef,
+    AssertTypeExpr, Block, MypyFile, FuncBase, FuncItem, CallExpr, ClassDef, Decorator, FuncDef,
     ExpressionStmt, AssignmentStmt, OperatorAssignmentStmt, WhileStmt,
     ForStmt, ReturnStmt, AssertStmt, DelStmt, IfStmt, RaiseStmt,
     TryStmt, WithStmt, MatchStmt, NameExpr, MemberExpr, OpExpr, SliceExpr, CastExpr,
@@ -205,6 +205,9 @@ def visit_slice_expr(self, o: SliceExpr) -> None:
     def visit_cast_expr(self, o: CastExpr) -> None:
         o.expr.accept(self)
 
+    def visit_assert_type_expr(self, o: AssertTypeExpr) -> None:
+        o.expr.accept(self)
+
     def visit_reveal_expr(self, o: RevealExpr) -> None:
         if o.kind == REVEAL_TYPE:
             assert o.expr is not None
diff --git a/mypy/treetransform.py b/mypy/treetransform.py
@@ -6,7 +6,7 @@
 from typing import List, Dict, cast, Optional, Iterable
 
 from mypy.nodes import (
-    MypyFile, Import, Node, ImportAll, ImportFrom, FuncItem, FuncDef,
+    AssertTypeExpr, MypyFile, Import, Node, ImportAll, ImportFrom, FuncItem, FuncDef,
     OverloadedFuncDef, ClassDef, Decorator, Block, Var,
     OperatorAssignmentStmt, ExpressionStmt, AssignmentStmt, ReturnStmt,
     RaiseStmt, AssertStmt, DelStmt, BreakStmt, ContinueStmt,
@@ -407,6 +407,9 @@ def visit_cast_expr(self, node: CastExpr) -> CastExpr:
         return CastExpr(self.expr(node.expr),
                         self.type(node.type))
 
+    def visit_assert_type_expr(self, node: AssertTypeExpr) -> AssertTypeExpr:
+        return AssertTypeExpr(self.expr(node.expr), self.type(node.type))
+
     def visit_reveal_expr(self, node: RevealExpr) -> RevealExpr:
         if node.kind == REVEAL_TYPE:
             assert node.expr is not None
diff --git a/mypy/types.py b/mypy/types.py
@@ -132,6 +132,11 @@
     'typing_extensions.reveal_type',
 )
 
+ASSERT_TYPE_NAMES: Final = (
+    'typing.assert_type',
+    'typing_extensions.assert_type',
+)
+
 # Attributes that can optionally be defined in the body of a subclass of
 # enum.Enum but are removed from the class __dict__ by EnumMeta.
 ENUM_REMOVED_PROPS: Final = (
diff --git a/mypy/visitor.py b/mypy/visitor.py
@@ -81,6 +81,10 @@ def visit_comparison_expr(self, o: 'mypy.nodes.ComparisonExpr') -> T:
     def visit_cast_expr(self, o: 'mypy.nodes.CastExpr') -> T:
         pass
 
+    @abstractmethod
+    def visit_assert_type_expr(self, o: 'mypy.nodes.AssertTypeExpr') -> T:
+        pass
+
     @abstractmethod
     def visit_reveal_expr(self, o: 'mypy.nodes.RevealExpr') -> T:
         pass
@@ -523,6 +527,9 @@ def visit_comparison_expr(self, o: 'mypy.nodes.ComparisonExpr') -> T:
     def visit_cast_expr(self, o: 'mypy.nodes.CastExpr') -> T:
         pass
 
+    def visit_assert_type_expr(self, o: 'mypy.nodes.AssertTypeExpr') -> T:
+        pass
+
     def visit_reveal_expr(self, o: 'mypy.nodes.RevealExpr') -> T:
         pass
 
diff --git a/mypyc/irbuild/expression.py b/mypyc/irbuild/expression.py
@@ -11,7 +11,7 @@
     ConditionalExpr, ComparisonExpr, IntExpr, FloatExpr, ComplexExpr, StrExpr,
     BytesExpr, EllipsisExpr, ListExpr, TupleExpr, DictExpr, SetExpr, ListComprehension,
     SetComprehension, DictionaryComprehension, SliceExpr, GeneratorExpr, CastExpr, StarExpr,
-    AssignmentExpr,
+    AssignmentExpr, AssertTypeExpr,
     Var, RefExpr, MypyFile, TypeInfo, TypeApplication, LDEF, ARG_POS
 )
 from mypy.types import TupleType, Instance, TypeType, ProperType, get_proper_type
@@ -203,6 +203,9 @@ def transform_super_expr(builder: IRBuilder, o: SuperExpr) -> Value:
 def transform_call_expr(builder: IRBuilder, expr: CallExpr) -> Value:
     if isinstance(expr.analyzed, CastExpr):
         return translate_cast_expr(builder, expr.analyzed)
+    elif isinstance(expr.analyzed, AssertTypeExpr):
+        # Compile to a no-op.
+        return builder.accept(expr.analyzed.expr)
 
     callee = expr.callee
     if isinstance(callee, IndexExpr) and isinstance(callee.analyzed, TypeApplication):
diff --git a/mypyc/irbuild/visitor.py b/mypyc/irbuild/visitor.py
@@ -6,7 +6,7 @@
 from typing_extensions import NoReturn
 
 from mypy.nodes import (
-    MypyFile, FuncDef, ReturnStmt, AssignmentStmt, OpExpr,
+    AssertTypeExpr, MypyFile, FuncDef, ReturnStmt, AssignmentStmt, OpExpr,
     IntExpr, NameExpr, Var, IfStmt, UnaryExpr, ComparisonExpr, WhileStmt, CallExpr,
     IndexExpr, Block, ListExpr, ExpressionStmt, MemberExpr, ForStmt,
     BreakStmt, ContinueStmt, ConditionalExpr, OperatorAssignmentStmt, TupleExpr, ClassDef,
@@ -327,6 +327,9 @@ def visit_var(self, o: Var) -> None:
     def visit_cast_expr(self, o: CastExpr) -> Value:
         assert False, "CastExpr should have been handled in CallExpr"
 
+    def visit_assert_type_expr(self, o: AssertTypeExpr) -> Value:
+        assert False, "AssertTypeExpr should have been handled in CallExpr"
+
     def visit_star_expr(self, o: StarExpr) -> Value:
         assert False, "should have been handled in Tuple/List/Set/DictExpr or CallExpr"
 
diff --git a/mypyc/test-data/irbuild-basic.test b/mypyc/test-data/irbuild-basic.test
@@ -876,6 +876,17 @@ L0:
     o = r3
     return 1
 
+[case testAssertType]
+from typing import assert_type
+def f(x: int) -> None:
+    y = assert_type(x, int)
+[out]
+def f(x):
+    x, y :: int
+L0:
+    y = x
+    return 1
+
 [case testDownCast]
 from typing import cast, List, Tuple
 class A: pass
diff --git a/test-data/unit/check-expressions.test b/test-data/unit/check-expressions.test
@@ -1036,6 +1036,18 @@ class B: pass
 [out]
 main:3: error: "A" not callable
 
+-- assert_type()
+
+[case testAssertType]
+from typing import assert_type, Any
+from typing_extensions import Literal
+a: int = 1
+returned = assert_type(a, int)
+reveal_type(returned)  # N: Revealed type is "builtins.int"
+assert_type(a, str)  # E: Expression is of type "int", not "str"
+assert_type(a, Any)  # E: Expression is of type "int", not "Any"
+assert_type(a, Literal[1])  # E: Expression is of type "int", not "Literal[1]"
+[builtins fixtures/tuple.pyi]
 
 -- None return type
 -- ----------------
diff --git a/test-data/unit/fixtures/typing-full.pyi b/test-data/unit/fixtures/typing-full.pyi
@@ -11,6 +11,7 @@ from abc import abstractmethod, ABCMeta
 class GenericMeta(type): pass
 
 def cast(t, o): ...
+def assert_type(o, t): ...
 overload = 0
 Any = 0
 Union = 0
diff --git a/test-data/unit/lib-stub/typing.pyi b/test-data/unit/lib-stub/typing.pyi
@@ -9,6 +9,7 @@
 # the stubs under fixtures/.
 
 cast = 0
+assert_type = 0
 overload = 0
 Any = 0
 Union = 0
diff --git a/test-data/unit/semanal-errors.test b/test-data/unit/semanal-errors.test
diff --git a/test-data/unit/semanal-types.test b/test-data/unit/semanal-types.test
diff --git a/test-data/unit/typexport-basic.test b/test-data/unit/typexport-basic.test

Original file line number	Diff line number	Diff line change
`@@ -113,6 +113,9 @@ def __str__(self) -> str:`
`113`	`113`	`REDUNDANT_CAST: Final = ErrorCode(`
`114`	`114`	`"redundant-cast", "Check that cast changes type of expression", "General"`
`115`	`115`	`)`
	`116`	`+ASSERT_TYPE: Final = ErrorCode(`
	`117`	`+ "assert-type", "Check that assert_type() call succeeds", "General"`
	`118`	`+)`
`116`	`119`	`COMPARISON_OVERLAP: Final = ErrorCode(`
`117`	`120`	`"comparison-overlap", "Check that types in comparisons and 'in' expressions overlap", "General"`
`118`	`121`	`)`