support typed dict NotRequired

Summary: This doesn't handle total=True/False yet, I'll do that later. Reviewed By: stroxler Differential Revision: D68292311 fbshipit-source-id: cf5c69b1af5e304943607ded39f396d2c201f7a2
facebook · Jan 17, 2025 · f7044fd · f7044fd
1 parent 3abc2d0
commit f7044fd
Show file tree

Hide file tree

Showing 6 changed files with 127 additions and 8 deletions.
diff --git a/pyre2/conformance/third_party/conformance.exp b/pyre2/conformance/third_party/conformance.exp
@@ -21437,6 +21437,46 @@
       "name": "PyreError",
       "stop_column": 15,
       "stop_line": 40
+    },
+    {
+      "code": -2,
+      "column": 14,
+      "concise_description": "EXPECTED C2 <: A2",
+      "description": "EXPECTED C2 <: A2",
+      "line": 79,
+      "name": "PyreError",
+      "stop_column": 15,
+      "stop_line": 79
+    },
+    {
+      "code": -2,
+      "column": 14,
+      "concise_description": "EXPECTED C2 <: B2",
+      "description": "EXPECTED C2 <: B2",
+      "line": 82,
+      "name": "PyreError",
+      "stop_column": 15,
+      "stop_line": 82
+    },
+    {
+      "code": -2,
+      "column": 14,
+      "concise_description": "EXPECTED A2 <: C2",
+      "description": "EXPECTED A2 <: C2",
+      "line": 84,
+      "name": "PyreError",
+      "stop_column": 15,
+      "stop_line": 84
+    },
+    {
+      "code": -2,
+      "column": 14,
+      "concise_description": "EXPECTED B2 <: C2",
+      "description": "EXPECTED B2 <: C2",
+      "line": 85,
+      "name": "PyreError",
+      "stop_column": 15,
+      "stop_line": 85
     }
   ],
   "typeddicts_readonly_inheritance.py": [
@@ -21646,6 +21686,46 @@
       "stop_column": 20,
       "stop_line": 49
     },
+    {
+      "code": -2,
+      "column": 7,
+      "concise_description": "EXPECTED TD4 <: TD3",
+      "description": "EXPECTED TD4 <: TD3",
+      "line": 53,
+      "name": "PyreError",
+      "stop_column": 10,
+      "stop_line": 53
+    },
+    {
+      "code": -2,
+      "column": 7,
+      "concise_description": "EXPECTED TD5 <: TD4",
+      "description": "EXPECTED TD5 <: TD4",
+      "line": 55,
+      "name": "PyreError",
+      "stop_column": 10,
+      "stop_line": 55
+    },
+    {
+      "code": -2,
+      "column": 7,
+      "concise_description": "EXPECTED TD5 <: TD4",
+      "description": "EXPECTED TD5 <: TD4",
+      "line": 56,
+      "name": "PyreError",
+      "stop_column": 10,
+      "stop_line": 56
+    },
+    {
+      "code": -2,
+      "column": 7,
+      "concise_description": "EXPECTED TD4 <: TD5",
+      "description": "EXPECTED TD4 <: TD5",
+      "line": 58,
+      "name": "PyreError",
+      "stop_column": 10,
+      "stop_line": 58
+    },
     {
       "code": -2,
       "column": 18,

diff --git a/pyre2/conformance/third_party/conformance.result b/pyre2/conformance/third_party/conformance.result
@@ -1558,9 +1558,7 @@
   "typeddicts_readonly_consistency.py": [
     "Line 38: Expected 1 errors",
     "Line 81: Expected 1 errors",
-    "Line 82: Expected 1 errors",
-    "Line 84: Expected 1 errors",
-    "Line 85: Expected 1 errors"
+    "Line 79: Unexpected errors ['EXPECTED C2 <: A2']"
   ],
   "typeddicts_readonly_inheritance.py": [
     "Line 50: Expected 1 errors",
@@ -1592,6 +1590,10 @@
     "Line 47: Unexpected errors ['EXPECTED dict[str, int] <: TD3']",
     "Line 48: Unexpected errors ['EXPECTED dict[str, int] <: TD4']",
     "Line 49: Unexpected errors ['EXPECTED dict[str, int] <: TD5']",
+    "Line 53: Unexpected errors ['EXPECTED TD4 <: TD3']",
+    "Line 55: Unexpected errors ['EXPECTED TD5 <: TD4']",
+    "Line 56: Unexpected errors ['EXPECTED TD5 <: TD4']",
+    "Line 58: Unexpected errors ['EXPECTED TD4 <: TD5']",
     "Line 68: Unexpected errors ['TODO: Answers::apply_special_form cannot handle `Required[int]`']",
     "Line 70: Unexpected errors ['TODO: Answers::apply_special_form cannot handle `Required[Annotated[int, \"\"]]`']",
     "Line 73: Unexpected errors ['Expected a callable, got type[TypedDict]']",

diff --git a/pyre2/conformance/third_party/results.json b/pyre2/conformance/third_party/results.json
@@ -3,7 +3,7 @@
   "pass": 9,
   "fail": 124,
   "pass_rate": 0.07,
-  "differences": 1364,
+  "differences": 1366,
   "passing": [
     "annotations_coroutines.py",
     "directives_no_type_check.py",
@@ -133,11 +133,11 @@
     "typeddicts_inheritance.py": 6,
     "typeddicts_operations.py": 7,
     "typeddicts_readonly.py": 8,
-    "typeddicts_readonly_consistency.py": 5,
+    "typeddicts_readonly_consistency.py": 3,
     "typeddicts_readonly_inheritance.py": 12,
     "typeddicts_readonly_kwargs.py": 1,
     "typeddicts_readonly_update.py": 3,
-    "typeddicts_required.py": 11,
+    "typeddicts_required.py": 15,
     "typeddicts_type_consistency.py": 15,
     "typeddicts_usage.py": 5
   },

diff --git a/pyre2/pyre2/bin/alt/classes.rs b/pyre2/pyre2/bin/alt/classes.rs
@@ -34,6 +34,7 @@ use crate::binding::binding::KeyLegacyTypeParam;
 use crate::graph::index::Idx;
 use crate::module::short_identifier::ShortIdentifier;
 use crate::types::annotation::Annotation;
+use crate::types::annotation::Qualifier;
 use crate::types::class::Class;
 use crate::types::class::ClassType;
 use crate::types::class::TArgs;
@@ -782,11 +783,28 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
             .iter()
             .filter_map(|name| {
                 if let Some(ClassField {
-                    annotation: Some(Annotation { ty: Some(ty), .. }),
+                    annotation:
+                        Some(Annotation {
+                            ty: Some(ty),
+                            qualifiers,
+                        }),
                     ..
                 }) = self.get_class_field(typed_dict.0.class_object(), name)
                 {
-                    Some((name.clone(), TypedDictField { ty, required: true }))
+                    Some((
+                        name.clone(),
+                        TypedDictField {
+                            ty,
+                            required: if qualifiers.contains(&Qualifier::Required) {
+                                true
+                            } else if qualifiers.contains(&Qualifier::NotRequired) {
+                                false
+                            } else {
+                                // TODO(yangdanny): consider total=True/False
+                                true
+                            },
+                        },
+                    ))
                 } else {
                     None
                 }

diff --git a/pyre2/pyre2/bin/test/stdlib.rs b/pyre2/pyre2/bin/test/stdlib.rs
@@ -103,6 +103,8 @@ Never: _SpecialForm
 NoReturn: _SpecialForm
 Annotated: _SpecialForm
 TypedDict: _SpecialForm
+Required: _SpecialForm
+NotRequired: _SpecialForm
 
 def assert_type(x, y) -> None: ...
 

diff --git a/pyre2/pyre2/bin/test/typed_dict.rs b/pyre2/pyre2/bin/test/typed_dict.rs
@@ -77,3 +77,20 @@ def foo(a: Coord, b: Coord3D, c: Pair):
     coord4: Pair = a
     "#,
 );
+
+testcase!(
+    test_typed_dict_not_required,
+    r#"
+from typing import TypedDict, NotRequired
+class Coord(TypedDict):
+    x: int
+    y: int
+class CoordNotRequired(TypedDict):
+    x: NotRequired[int]
+    y: NotRequired[int]
+
+def foo(a: Coord, b: CoordNotRequired):
+    coord: Coord = b  # E: EXPECTED CoordNotRequired <: Coord
+    coord2: CoordNotRequired = a  # E: EXPECTED Coord <: CoordNotRequired
+    "#,
+);