Support forward references in assert_type

Reviewed By: samwgoldman Differential Revision: D70024022 fbshipit-source-id: ab0bbc1b5b608d371edd061978fa3316046efe0c
facebook · Feb 25, 2025 · 542fed5 · 542fed5
1 parent b349638
commit 542fed5
Show file tree

Hide file tree

Showing 6 changed files with 43 additions and 29 deletions.
diff --git a/pyre2/conformance/third_party/conformance.exp b/pyre2/conformance/third_party/conformance.exp
@@ -4877,26 +4877,6 @@
     }
   ],
   "directives_assert_type.py": [
-    {
-      "code": -2,
-      "column": 5,
-      "concise_description": "assert_type(ForwardReference, Any) failed",
-      "description": "assert_type(ForwardReference, Any) failed",
-      "line": 24,
-      "name": "PyreError",
-      "stop_column": 39,
-      "stop_line": 24
-    },
-    {
-      "code": -2,
-      "column": 20,
-      "concise_description": "untype, got Literal['ForwardReference']",
-      "description": "untype, got Literal['ForwardReference']",
-      "line": 24,
-      "name": "PyreError",
-      "stop_column": 38,
-      "stop_line": 24
-    },
     {
       "code": -2,
       "column": 5,

diff --git a/pyre2/conformance/third_party/conformance.result b/pyre2/conformance/third_party/conformance.result
@@ -376,9 +376,7 @@
     "Line 106: Unexpected errors ['assert_type(BoundMethod[DC6, (str) -> int], (str) -> int) failed']",
     "Line 205: Unexpected errors ['EXPECTED DC15 <: DataclassProto']"
   ],
-  "directives_assert_type.py": [
-    "Line 24: Unexpected errors ['assert_type(ForwardReference, Any) failed', \"untype, got Literal['ForwardReference']\"]"
-  ],
+  "directives_assert_type.py": [],
   "directives_cast.py": [
     "Line 16: Expected 1 errors"
   ],

diff --git a/pyre2/conformance/third_party/results.json b/pyre2/conformance/third_party/results.json
@@ -1,13 +1,14 @@
 {
   "total": 133,
-  "pass": 21,
-  "fail": 112,
-  "pass_rate": 0.16,
-  "differences": 807,
+  "pass": 22,
+  "fail": 111,
+  "pass_rate": 0.17,
+  "differences": 806,
   "passing": [
     "annotations_coroutines.py",
     "dataclasses_hash.py",
     "dataclasses_order.py",
+    "directives_assert_type.py",
     "directives_no_type_check.py",
     "directives_type_ignore.py",
     "directives_type_ignore_file2.py",
@@ -64,7 +65,6 @@
     "dataclasses_transform_func.py": 7,
     "dataclasses_transform_meta.py": 9,
     "dataclasses_usage.py": 7,
-    "directives_assert_type.py": 1,
     "directives_cast.py": 1,
     "directives_deprecated.py": 11,
     "directives_reveal_type.py": 4,

diff --git a/pyre2/pyre2/lib/binding/expr.rs b/pyre2/pyre2/lib/binding/expr.rs
@@ -13,6 +13,7 @@ use ruff_python_ast::Comprehension;
 use ruff_python_ast::Decorator;
 use ruff_python_ast::Expr;
 use ruff_python_ast::ExprBoolOp;
+use ruff_python_ast::ExprCall;
 use ruff_python_ast::ExprLambda;
 use ruff_python_ast::ExprSubscript;
 use ruff_python_ast::Identifier;
@@ -141,6 +142,27 @@ impl<'a> BindingsBuilder<'a> {
                 self.negate_and_merge_flow(base, &narrow_ops, None, *range);
                 return;
             }
+            Expr::Call(ExprCall {
+                range: _,
+                func,
+                arguments,
+            }) if self.as_special_export(func) == Some(SpecialExport::AssertType)
+                && arguments.args.len() > 1 =>
+            {
+                // Handle forward references in the second argument to an assert_type call
+                self.ensure_expr(func);
+                for (i, arg) in arguments.args.iter_mut().enumerate() {
+                    if i == 1 {
+                        self.ensure_type(arg, &mut None);
+                    } else {
+                        self.ensure_expr(arg);
+                    }
+                }
+                for kw in arguments.keywords.iter_mut() {
+                    self.ensure_expr(&mut kw.value);
+                }
+                return;
+            }
             Expr::Name(x) => {
                 let name = Ast::expr_name_identifier(x.clone());
                 let binding = self.forward_lookup(&name);

diff --git a/pyre2/pyre2/lib/export/special.rs b/pyre2/pyre2/lib/export/special.rs
@@ -27,6 +27,7 @@ pub enum SpecialExport {
     TypedDict,
     CollectionsNamedTuple,
     TypingNamedTuple,
+    AssertType,
 }
 
 impl SpecialExport {
@@ -44,6 +45,7 @@ impl SpecialExport {
             "TypedDict" => Some(Self::TypedDict),
             "namedtuple" => Some(Self::CollectionsNamedTuple),
             "NamedTuple" => Some(Self::TypingNamedTuple),
+            "assert_type" => Some(Self::AssertType),
             _ => None,
         }
     }
@@ -57,7 +59,8 @@ impl SpecialExport {
             | Self::Annotated
             | Self::Literal
             | Self::TypedDict
-            | Self::TypingNamedTuple => {
+            | Self::TypingNamedTuple
+            | Self::AssertType => {
                 matches!(m.as_str(), "typing" | "typing_extensions")
             }
             Self::CollectionsNamedTuple => matches!(m.as_str(), "collections"),

diff --git a/pyre2/pyre2/lib/test/simple.rs b/pyre2/pyre2/lib/test/simple.rs
@@ -1165,3 +1165,14 @@ from typing import Literal as L
 x: L["foo"] = "foo"
 "#,
 );
+
+testcase!(
+    test_assert_type_forward_ref,
+    r#"
+from typing import assert_type
+x: "ForwardRef"
+assert_type(x, "ForwardRef")
+class ForwardRef:
+    pass
+    "#,
+);