red-knot: infer string literal types

crates/red_knot_python_semantic/src/types.rs

@@ -181,6 +181,8 @@ pub enum Type<'db> {
     IntLiteral(i64),
     /// A boolean literal, either `True` or `False`.
     BooleanLiteral(bool),
+    /// A string literal
+    StringLiteral(StringLiteralType<'db>),
     /// A bytes literal
     BytesLiteral(BytesLiteralType<'db>),
     // TODO protocols, callable types, overloads, generics, type vars
@@ -278,6 +280,11 @@ impl<'db> Type<'db> {
                 Type::Unknown
             }
             Type::BooleanLiteral(_) => Type::Unknown,
+            Type::StringLiteral(_) => {
+                // TODO fix this comment with whatever it is we should say we
+                // want to do in the future.
+                Type::Unknown
+            }
             Type::BytesLiteral(_) => {
                 // TODO defer to Type::Instance(<bytes from typeshed>).member
                 Type::Unknown
@@ -378,6 +385,12 @@ pub struct IntersectionType<'db> {
     negative: FxOrderSet<Type<'db>>,
 }
 
+#[salsa::interned]
+pub struct StringLiteralType<'db> {
+    #[return_ref]
+    value: String,
+}
+
 #[salsa::interned]
 pub struct BytesLiteralType<'db> {
     #[return_ref]

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;
+use ruff_python_literal::escape::{AsciiEscape, Escape};
 
 use crate::types::{IntersectionType, Type, UnionType};
 use crate::Db;
@@ -49,6 +49,16 @@ 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("]")
+            }
         }
     }
 }

crates/red_knot_python_semantic/src/types/infer.rs

@@ -44,7 +44,7 @@ use crate::semantic_index::SemanticIndex;
 use crate::types::diagnostic::{TypeCheckDiagnostic, TypeCheckDiagnostics};
 use crate::types::{
     builtins_symbol_ty_by_name, definitions_ty, global_symbol_ty_by_name, BytesLiteralType,
-    ClassType, FunctionType, Name, Type, UnionBuilder,
+    ClassType, FunctionType, Name, StringLiteralType, Type, UnionBuilder,
 };
 use crate::Db;
 
@@ -1243,9 +1243,8 @@ impl<'db> TypeInferenceBuilder<'db> {
     }
 
     #[allow(clippy::unused_self)]
-    fn infer_string_literal_expression(&mut self, _literal: &ast::ExprStringLiteral) -> Type<'db> {
-        // TODO Literal["..."] or str
-        Type::Unknown
+    fn infer_string_literal_expression(&mut self, literal: &ast::ExprStringLiteral) -> Type<'db> {
+        Type::StringLiteral(StringLiteralType::new(self.db, literal.value.to_string()))
     }
 
     #[allow(clippy::unused_self)]
@@ -1785,6 +1784,17 @@ impl<'db> TypeInferenceBuilder<'db> {
                     _ => Type::Unknown, // TODO
                 }
             }
+            Type::StringLiteral(lhs) => match right_ty {
+                Type::StringLiteral(rhs) => match op {
+                    ast::Operator::Add => Type::StringLiteral(StringLiteralType::new(self.db, {
+                        let lhs_value = lhs.value(self.db);
+                        let rhs_value = rhs.value(self.db);
+                        lhs_value.clone() + rhs_value
+                    })),
+                    _ => Type::Unknown, // TODO
+                },
+                _ => Type::Unknown, // TODO
+            },
             _ => Type::Unknown, // TODO
         }
     }
@@ -2298,6 +2308,28 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn string_type() -> anyhow::Result<()> {
+        let mut db = setup_db();
+
+        db.write_dedented(
+            "src/a.py",
+            r#"
+            w = "Hello"
+            x = 'world'
+            y = "Guten " + 'tag'
+            z = 'bon ' + "jour"
+            "#,
+        )?;
+
+        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]");
+
+        Ok(())
+    }
+
     #[test]
     fn bytes_type() -> anyhow::Result<()> {
         let mut db = setup_db();