Fix TypedDict

Author: charliermarsh

crates/ty_python_semantic/resources/mdtest/typed_dict.md

@@ -71,6 +71,87 @@ takes_formatter({"format": "%(message)s"})
 takes_formatter({"factory": object(), "facility": "local0"})
 ```
 
+Large `dict[str, TypedDict]` literals should still preserve inner `TypedDict` keys after crossing
+the large-collection promotion threshold:
+
+```py
+from typing import TypedDict
+
+class Entry(TypedDict):
+    a: str
+    b: bool
+
+entries: dict[str, Entry] = {
+    "k0": {"a": "v", "b": False},
+    "k1": {"a": "v", "b": False},
+    "k2": {"a": "v", "b": False},
+    "k3": {"a": "v", "b": False},
+    "k4": {"a": "v", "b": False},
+    "k5": {"a": "v", "b": False},
+    "k6": {"a": "v", "b": False},
+    "k7": {"a": "v", "b": False},
+    "k8": {"a": "v", "b": False},
+    "k9": {"a": "v", "b": False},
+    "k10": {"a": "v", "b": False},
+    "k11": {"a": "v", "b": False},
+    "k12": {"a": "v", "b": False},
+    "k13": {"a": "v", "b": False},
+    "k14": {"a": "v", "b": False},
+    "k15": {"a": "v", "b": False},
+    "k16": {"a": "v", "b": False},
+    "k17": {"a": "v", "b": False},
+    "k18": {"a": "v", "b": False},
+    "k19": {"a": "v", "b": False},
+    "k20": {"a": "v", "b": False},
+    "k21": {"a": "v", "b": False},
+    "k22": {"a": "v", "b": False},
+    "k23": {"a": "v", "b": False},
+    "k24": {"a": "v", "b": False},
+    "k25": {"a": "v", "b": False},
+    "k26": {"a": "v", "b": False},
+    "k27": {"a": "v", "b": False},
+    "k28": {"a": "v", "b": False},
+    "k29": {"a": "v", "b": False},
+    "k30": {"a": "v", "b": False},
+    "k31": {"a": "v", "b": False},
+    "k32": {"a": "v", "b": False},
+    "k33": {"a": "v", "b": False},
+    "k34": {"a": "v", "b": False},
+    "k35": {"a": "v", "b": False},
+    "k36": {"a": "v", "b": False},
+    "k37": {"a": "v", "b": False},
+    "k38": {"a": "v", "b": False},
+    "k39": {"a": "v", "b": False},
+    "k40": {"a": "v", "b": False},
+    "k41": {"a": "v", "b": False},
+    "k42": {"a": "v", "b": False},
+    "k43": {"a": "v", "b": False},
+    "k44": {"a": "v", "b": False},
+    "k45": {"a": "v", "b": False},
+    "k46": {"a": "v", "b": False},
+    "k47": {"a": "v", "b": False},
+    "k48": {"a": "v", "b": False},
+    "k49": {"a": "v", "b": False},
+    "k50": {"a": "v", "b": False},
+    "k51": {"a": "v", "b": False},
+    "k52": {"a": "v", "b": False},
+    "k53": {"a": "v", "b": False},
+    "k54": {"a": "v", "b": False},
+    "k55": {"a": "v", "b": False},
+    "k56": {"a": "v", "b": False},
+    "k57": {"a": "v", "b": False},
+    "k58": {"a": "v", "b": False},
+    "k59": {"a": "v", "b": False},
+    "k60": {"a": "v", "b": False},
+    "k61": {"a": "v", "b": False},
+    "k62": {"a": "v", "b": False},
+    "k63": {"a": "v", "b": False},
+    "k64": {"a": "v", "b": False},
+}
+
+reveal_type(entries["k0"])  # revealed: Entry
+```
+
 Methods that are available on `dict`s are also available on `TypedDict`s:
 
 ```py

crates/ty_python_semantic/src/types/infer/builder/typed_dict.rs

@@ -11,6 +11,7 @@ use crate::types::diagnostic::{
     INVALID_ARGUMENT_TYPE, INVALID_TYPE_FORM, MISSING_ARGUMENT, TOO_MANY_POSITIONAL_ARGUMENTS,
     UNKNOWN_ARGUMENT, report_mismatched_type_name,
 };
+use crate::types::infer::InferenceFlags;
 use crate::types::infer::builder::DeferredExpressionState;
 use crate::types::special_form::TypeQualifier;
 use crate::types::typed_dict::{
@@ -61,6 +62,25 @@ impl<'expr> TypedDictConstructorForm<'expr> {
 }
 
 impl<'db> TypeInferenceBuilder<'db, '_> {
+    /// Preserve string literal keys while inferring `TypedDict` fields, even when the enclosing
+    /// collection enables large-literal promotion.
+    fn infer_typed_dict_key_expression(&mut self, key: &ast::Expr) -> Type<'db> {
+        if let Some(key_ty) = self.try_expression_type(key) {
+            return key_ty.as_string_literal().map_or_else(
+                || {
+                    key.as_string_literal_expr().map_or(key_ty, |literal| {
+                        Type::string_literal(self.db(), literal.value.to_str())
+                    })
+                },
+                |_| key_ty,
+            );
+        }
+
+        self.with_inference_flag(InferenceFlags::PROMOTE_LITERALS, false, |builder| {
+            builder.infer_expression(key, TypeContext::default())
+        })
+    }
+
     /// Infer a `TypedDict(name, fields)` call expression.
     ///
     /// This method *does not* call `infer_expression` on the object being called;
@@ -309,7 +329,10 @@ impl<'db> TypeInferenceBuilder<'db, '_> {
         let typed_dict_items = typed_dict.items(self.db());
 
         for item in items {
-            let key_ty = self.infer_optional_expression(item.key.as_ref(), TypeContext::default());
+            let key_ty = item
+                .key
+                .as_ref()
+                .map(|key| self.infer_typed_dict_key_expression(key));
             if let Some((key, key_ty)) = item.key.as_ref().zip(key_ty) {
                 item_types.insert(key.node_index().load(), key_ty);
             }
@@ -436,7 +459,7 @@ impl<'db> TypeInferenceBuilder<'db, '_> {
             let value_tcx = item
                 .key
                 .as_ref()
-                .map(|key| self.get_or_infer_expression(key, TypeContext::default()))
+                .map(|key| self.infer_typed_dict_key_expression(key))
                 .and_then(Type::as_string_literal)
                 .and_then(|key| items.get(key.value(self.db())))
                 .map(|field| TypeContext::new(Some(field.declared_ty)))