From 594b7b04d3b04bcf42861f86207017c8117678ca Mon Sep 17 00:00:00 2001
From: Alex Waygood <Alex.Waygood@Gmail.com>
Date: Fri, 28 Nov 2025 18:40:34 +0000
Subject: [PATCH] [ty] Preserve quoting style when autofixing `TypedDict` keys
 (#21682)

---
 ...ict`_-_Diagnostics_(e5289abf5c570c29).snap | 26 +++++++++++++++++++
 .../resources/mdtest/typed_dict.md            |  8 ++++++
 .../src/types/diagnostic.rs                   | 12 ++++++---
 3 files changed, 42 insertions(+), 4 deletions(-)

diff --git a/crates/ty_python_semantic/resources/mdtest/snapshots/typed_dict.md_-_`TypedDict`_-_Diagnostics_(e5289abf5c570c29).snap b/crates/ty_python_semantic/resources/mdtest/snapshots/typed_dict.md_-_`TypedDict`_-_Diagnostics_(e5289abf5c570c29).snap
index 0fc0bbd2a6..1f70b8acd3 100644
--- a/crates/ty_python_semantic/resources/mdtest/snapshots/typed_dict.md_-_`TypedDict`_-_Diagnostics_(e5289abf5c570c29).snap
+++ b/crates/ty_python_semantic/resources/mdtest/snapshots/typed_dict.md_-_`TypedDict`_-_Diagnostics_(e5289abf5c570c29).snap
@@ -52,6 +52,9 @@ mdtest path: crates/ty_python_semantic/resources/mdtest/typed_dict.md
 38 | 
 39 | def write_to_readonly_key(employee: Employee):
 40 |     employee["id"] = 42  # error: [invalid-assignment]
+41 | def write_to_non_existing_key_single_quotes(person: Person):
+42 |     # error: [invalid-key]
+43 |     person['naem'] = "Alice"  # fmt: skip
 ```
 
 # Diagnostics
@@ -217,6 +220,8 @@ error[invalid-assignment]: Cannot assign to key "id" on TypedDict `Employee`
    |     -------- ^^^^ key is marked read-only
    |     |
    |     TypedDict `Employee`
+41 | def write_to_non_existing_key_single_quotes(person: Person):
+42 |     # error: [invalid-key]
    |
 info: Item declaration
   --> src/mdtest_snippet.py:36:5
@@ -229,3 +234,24 @@ info: Item declaration
 info: rule `invalid-assignment` is enabled by default
 
 ```
+
+```
+error[invalid-key]: Unknown key "naem" for TypedDict `Person`
+  --> src/mdtest_snippet.py:43:5
+   |
+41 | def write_to_non_existing_key_single_quotes(person: Person):
+42 |     # error: [invalid-key]
+43 |     person['naem'] = "Alice"  # fmt: skip
+   |     ------ ^^^^^^ Did you mean 'name'?
+   |     |
+   |     TypedDict `Person`
+   |
+info: rule `invalid-key` is enabled by default
+40 |     employee["id"] = 42  # error: [invalid-assignment]
+41 | def write_to_non_existing_key_single_quotes(person: Person):
+42 |     # error: [invalid-key]
+   -     person['naem'] = "Alice"  # fmt: skip
+43 +     person['name'] = "Alice"  # fmt: skip
+note: This is an unsafe fix and may change runtime behavior
+
+```
diff --git a/crates/ty_python_semantic/resources/mdtest/typed_dict.md b/crates/ty_python_semantic/resources/mdtest/typed_dict.md
index c843150d41..aadf8249ae 100644
--- a/crates/ty_python_semantic/resources/mdtest/typed_dict.md
+++ b/crates/ty_python_semantic/resources/mdtest/typed_dict.md
@@ -1502,6 +1502,14 @@ def write_to_readonly_key(employee: Employee):
     employee["id"] = 42  # error: [invalid-assignment]
 ```
 
+If the key uses single quotes, the autofix preserves that quoting style:
+
+```py
+def write_to_non_existing_key_single_quotes(person: Person):
+    # error: [invalid-key]
+    person['naem'] = "Alice"  # fmt: skip
+```
+
 ## Import aliases
 
 `TypedDict` can be imported with aliases and should work correctly:
diff --git a/crates/ty_python_semantic/src/types/diagnostic.rs b/crates/ty_python_semantic/src/types/diagnostic.rs
index 1266e62ef2..660af238bb 100644
--- a/crates/ty_python_semantic/src/types/diagnostic.rs
+++ b/crates/ty_python_semantic/src/types/diagnostic.rs
@@ -40,7 +40,7 @@ use ruff_db::{
 use ruff_diagnostics::{Edit, Fix};
 use ruff_python_ast::name::Name;
 use ruff_python_ast::parenthesize::parentheses_iterator;
-use ruff_python_ast::{self as ast, AnyNodeRef};
+use ruff_python_ast::{self as ast, AnyNodeRef, StringFlags};
 use ruff_python_trivia::CommentRanges;
 use ruff_text_size::{Ranged, TextRange};
 use rustc_hash::FxHashSet;
@@ -3456,11 +3456,15 @@ pub(crate) fn report_invalid_key_on_typed_dict<'db>(
 
                 let existing_keys = items.keys();
                 if let Some(suggestion) = did_you_mean(existing_keys, key) {
-                    if key_node.is_expr_string_literal() {
+                    if let AnyNodeRef::ExprStringLiteral(literal) = key_node {
+                        let quoted_suggestion = format!(
+                            "{quote}{suggestion}{quote}",
+                            quote = literal.value.first_literal_flags().quote_str()
+                        );
                         diagnostic
-                            .set_primary_message(format_args!("Did you mean \"{suggestion}\"?"));
+                            .set_primary_message(format_args!("Did you mean {quoted_suggestion}?"));
                         diagnostic.set_fix(Fix::unsafe_edit(Edit::range_replacement(
-                            format!("\"{suggestion}\""),
+                            quoted_suggestion,
                             key_node.range(),
                         )));
                     } else {