red-knot: properly quote inferred string literals

astral-sh · Aug 26, 2024 · 0195e27 · 0195e27
1 parent d9c8743
commit 0195e27
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 15 deletions.
diff --git a/crates/red_knot_python_semantic/src/types/display.rs b/crates/red_knot_python_semantic/src/types/display.rs
@@ -3,7 +3,7 @@
 use std::fmt::{Display, Formatter};
 
 use ruff_python_ast::str::Quote;
-use ruff_python_literal::escape::{AsciiEscape, Escape};
+use ruff_python_literal::escape::AsciiEscape;
 
 use crate::types::{IntersectionType, Type, UnionType};
 use crate::Db;
@@ -41,6 +41,11 @@ impl Display for DisplayType<'_> {
             Type::BooleanLiteral(boolean) => {
                 write!(f, "Literal[{}]", if *boolean { "True" } else { "False" })
             }
+            Type::StringLiteral(string) => write!(
+                f,
+                r#"Literal["{}"]"#,
+                string.value(self.db).replace('"', r#"\""#)
+            ),
             Type::BytesLiteral(bytes) => {
                 let escape =
                     AsciiEscape::with_preferred_quote(bytes.value(self.db).as_ref(), Quote::Double);
@@ -49,16 +54,6 @@ impl Display for DisplayType<'_> {
                 escape.bytes_repr().write(f)?;
                 f.write_str("]")
             }
-            Type::StringLiteral(string) => {
-                let escape = AsciiEscape::with_preferred_quote(
-                    string.value(self.db).as_bytes(),
-                    Quote::Double,
-                );
-
-                f.write_str("Literal[")?;
-                escape.write_body(f)?;
-                f.write_str("]")
-            }
         }
     }
 }

diff --git a/crates/red_knot_python_semantic/src/types/infer.rs b/crates/red_knot_python_semantic/src/types/infer.rs
@@ -2322,10 +2322,30 @@ mod tests {
             "#,
         )?;
 
-        assert_public_ty(&db, "src/a.py", "w", "Literal[Hello]");
-        assert_public_ty(&db, "src/a.py", "x", "Literal[world]");
-        assert_public_ty(&db, "src/a.py", "y", "Literal[Guten tag]");
-        assert_public_ty(&db, "src/a.py", "z", "Literal[bon jour]");
+        assert_public_ty(&db, "src/a.py", "w", r#"Literal["Hello"]"#);
+        assert_public_ty(&db, "src/a.py", "x", r#"Literal["world"]"#);
+        assert_public_ty(&db, "src/a.py", "y", r#"Literal["Guten tag"]"#);
+        assert_public_ty(&db, "src/a.py", "z", r#"Literal["bon jour"]"#);
+
+        Ok(())
+    }
+
+    #[test]
+    fn string_type_with_nested_quotes() -> anyhow::Result<()> {
+        let mut db = setup_db();
+
+        db.write_dedented(
+            "src/a.py",
+            r#"
+            x = 'I say "hello" to you'
+            y = "You say \"hey\" back"
+            z = 'No "closure here'
+            "#,
+        )?;
+
+        assert_public_ty(&db, "src/a.py", "x", r#"Literal["I say \"hello\" to you"]"#);
+        assert_public_ty(&db, "src/a.py", "y", r#"Literal["You say \"hey\" back"]"#);
+        assert_public_ty(&db, "src/a.py", "z", r#"Literal["No \"closure here"]"#);
 
         Ok(())
     }