[ty] Ensuring TypedDict subscripts for unknown keys return Unknown

Author: charliermarsh

crates/ty_python_semantic/resources/mdtest/typed_dict.md

@@ -1414,6 +1414,7 @@ AGE_FINAL: Final[Literal["age"]] = "age"
 
 def _(
     person: Person,
+    animal: Animal,
     being: Person | Animal,
     literal_key: Literal["age"],
     union_of_keys: Literal["age", "name"],
@@ -1439,12 +1440,13 @@ def _(
     # No error here:
     reveal_type(person[unknown_key])  # revealed: Unknown
 
+    # error: [invalid-key] "Unknown key "anything" for TypedDict `Animal`"
+    reveal_type(animal["anything"])  # revealed: Unknown
+
     reveal_type(being["name"])  # revealed: str
 
-    # TODO: A type of `int | None | Unknown` might be better here. The `str` is mixed in
-    # because `Animal.__getitem__` can only return `str`.
     # error: [invalid-key] "Unknown key "age" for TypedDict `Animal`"
-    reveal_type(being["age"])  # revealed: int | None | str
+    reveal_type(being["age"])  # revealed: int | None | Unknown
 ```
 
 ### Writing

crates/ty_python_semantic/src/types/subscript.rs

@@ -23,7 +23,7 @@ use super::special_form::SpecialFormType;
 use super::tuple::TupleSpec;
 use super::{
     DynamicType, IntersectionBuilder, IntersectionType, KnownInstanceType, Type, TypeAliasType,
-    UnionBuilder, UnionType, todo_type,
+    TypedDictType, UnionBuilder, UnionType, todo_type,
 };
 
 /// The kind of subscriptable type that had an out-of-bounds index.
@@ -120,6 +120,11 @@ pub(crate) enum SubscriptErrorKind<'db> {
         kind: CallErrorKind,
         bindings: Box<Bindings<'db>>,
     },
+    /// A `TypedDict` was subscripted with an invalid key.
+    InvalidTypedDictKey {
+        typed_dict: TypedDictType<'db>,
+        slice_ty: Type<'db>,
+    },
     /// The type does not support subscripting via the expected dunder.
     NotSubscriptable {
         value_ty: Type<'db>,
@@ -280,6 +285,21 @@ impl<'db> SubscriptErrorKind<'db> {
                     }
                 }
             },
+            Self::InvalidTypedDictKey {
+                typed_dict,
+                slice_ty,
+            } => {
+                let typed_dict_ty = Type::TypedDict(*typed_dict);
+                report_invalid_key_on_typed_dict(
+                    context,
+                    value_node.into(),
+                    slice_node.into(),
+                    typed_dict_ty,
+                    None,
+                    *slice_ty,
+                    typed_dict.items(db),
+                );
+            }
             Self::NotSubscriptable { value_ty, method } => {
                 report_not_subscriptable(context, subscript, *value_ty, method.as_str());
             }
@@ -411,6 +431,45 @@ where
     ))
 }
 
+// `TypedDict` subscripts need custom handling because invalid keys should still
+// recover with `Unknown` while emitting `invalid-key`, which is not naturally
+// representable via synthesized `__getitem__` overloads alone.
+fn typed_dict_subscript<'db>(
+    db: &'db dyn Db,
+    typed_dict: TypedDictType<'db>,
+    slice_ty: Type<'db>,
+) -> Result<Type<'db>, SubscriptError<'db>> {
+    if slice_ty.is_dynamic() {
+        return Ok(Type::unknown());
+    }
+
+    let Some(key) = slice_ty
+        .as_string_literal()
+        .map(|literal| literal.value(db))
+    else {
+        return Err(SubscriptError::new(
+            Type::unknown(),
+            SubscriptErrorKind::InvalidTypedDictKey {
+                typed_dict,
+                slice_ty,
+            },
+        ));
+    };
+
+    typed_dict.items(db).get(key).map_or_else(
+        || {
+            Err(SubscriptError::new(
+                Type::unknown(),
+                SubscriptErrorKind::InvalidTypedDictKey {
+                    typed_dict,
+                    slice_ty,
+                },
+            ))
+        },
+        |field| Ok(field.declared_ty),
+    )
+}
+
 impl<'db> Type<'db> {
     pub(super) fn subscript(
         self,
@@ -451,6 +510,11 @@ impl<'db> Type<'db> {
                 }))
             }
 
+            // Ex) Given `person["name"]`, return `str`
+            (Type::TypedDict(typed_dict), _) if expr_context != ast::ExprContext::Store => {
+                Some(typed_dict_subscript(db, typed_dict, slice_ty))
+            }
+
             // Ex) Given `("a", "b", "c", "d")[1]`, return `"b"`
             (Type::NominalInstance(nominal), Type::LiteralValue(literal)) if literal.is_int() => {
                 let i64_int = literal.as_int().unwrap();