[ty] allow any string Literal type expression as a key when constructing a TypedDict (#20792)

Author: mtshiba

crates/ty_python_semantic/resources/mdtest/typed_dict.md

@@ -46,6 +46,46 @@ Methods that are available on `dict`s are also available on `TypedDict`s:
 bob.update(age=26)
 ```
 
+`TypedDict` keys do not have to be string literals, as long as they can be statically determined
+(inferred to be of type string `Literal`).
+
+```py
+from typing import Literal, Final
+
+NAME = "name"
+AGE = "age"
+
+def non_literal() -> str:
+    return "name"
+
+def name_or_age() -> Literal["name", "age"]:
+    return "name"
+
+carol: Person = {NAME: "Carol", AGE: 20}
+
+reveal_type(carol[NAME])  # revealed: str
+# error: [invalid-key] "TypedDict `Person` cannot be indexed with a key of type `str`"
+reveal_type(carol[non_literal()])  # revealed: Unknown
+reveal_type(carol[name_or_age()])  # revealed: str | int | None
+
+FINAL_NAME: Final = "name"
+FINAL_AGE: Final = "age"
+
+def _():
+    carol: Person = {FINAL_NAME: "Carol", FINAL_AGE: 20}
+
+CAPITALIZED_NAME = "Name"
+
+# error: [invalid-key] "Invalid key access on TypedDict `Person`: Unknown key "Name" - did you mean "name"?"
+# error: [missing-typed-dict-key] "Missing required key 'name' in TypedDict `Person` constructor"
+dave: Person = {CAPITALIZED_NAME: "Dave", "age": 20}
+
+def age() -> Literal["age"] | None:
+    return "age"
+
+eve: Person = {"na" + "me": "Eve", age() or "age": 20}
+```
+
 The construction of a `TypedDict` is checked for type correctness:
 
 ```py

crates/ty_python_semantic/src/types/typed_dict.rs

@@ -387,21 +387,22 @@ fn validate_from_keywords<'db, 'ast>(
 
 /// Validates a `TypedDict` dictionary literal assignment,
 /// e.g. `person: Person = {"name": "Alice", "age": 30}`
-pub(super) fn validate_typed_dict_dict_literal<'db, 'ast>(
-    context: &InferContext<'db, 'ast>,
+pub(super) fn validate_typed_dict_dict_literal<'db>(
+    context: &InferContext<'db, '_>,
     typed_dict: TypedDictType<'db>,
-    dict_expr: &'ast ast::ExprDict,
-    error_node: AnyNodeRef<'ast>,
+    dict_expr: &ast::ExprDict,
+    error_node: AnyNodeRef,
     expression_type_fn: impl Fn(&ast::Expr) -> Type<'db>,
-) -> Result<OrderSet<&'ast str>, OrderSet<&'ast str>> {
+) -> Result<OrderSet<&'db str>, OrderSet<&'db str>> {
     let mut valid = true;
     let mut provided_keys = OrderSet::new();
 
     // Validate each key-value pair in the dictionary literal
     for item in &dict_expr.items {
         if let Some(key_expr) = &item.key {
-            if let ast::Expr::StringLiteral(key_literal) = key_expr {
-                let key_str = key_literal.value.to_str();
+            let key_ty = expression_type_fn(key_expr);
+            if let Type::StringLiteral(key_str) = key_ty {
+                let key_str = key_str.value(context.db());
                 provided_keys.insert(key_str);
 
                 let value_type = expression_type_fn(&item.value);