[red-knot] support fstring expressions

crates/red_knot_python_semantic/src/types.rs

@@ -721,6 +721,45 @@ impl<'db> Type<'db> {
             Type::Tuple(_) => builtins_symbol_ty(db, "tuple"),
         }
     }
+
+    /// Return the string representation of this type when converted to string as it would be
+    /// provided by the `__str__` method. If that can't be determined, return `None`.
+    ///
+    /// When not available, this should fall back to the value of `[Type::repr]`.
+    /// Note: this method is used in the builtins `format`, `print`, `str.format` and `f-strings`.
+    pub fn str(&self, db: &'db dyn Db) -> Option<Type<'db>> {
+        let str_result = match self {
+            Type::IntLiteral(_) => None,
+            Type::BooleanLiteral(_) => None,
+            Type::StringLiteral(_) => Some(*self),
+            // TODO: handle more complex types
+            _ => None,
+        };
+        str_result.or_else(|| self.repr(db))
+    }
+
+    /// Return the string representation of this type as it would be provided by the  `__repr__`
+    /// method at runtime. If that can't be determined, return `None`.
+    pub fn repr(&self, db: &'db dyn Db) -> Option<Type<'db>> {
+        match self {
+            Type::IntLiteral(number) => Some(Type::StringLiteral(StringLiteralType::new(db, {
+                number.to_string().into_boxed_str()
+            }))),
+            Type::BooleanLiteral(true) => Some(Type::StringLiteral(StringLiteralType::new(db, {
+                "True".into()
+            }))),
+            Type::BooleanLiteral(false) => Some(Type::StringLiteral(StringLiteralType::new(db, {
+                "False".into()
+            }))),
+            Type::StringLiteral(literal) => {
+                Some(Type::StringLiteral(StringLiteralType::new(db, {
+                    format!("'{}'", literal.value(db)).into()
+                })))
+            }
+            // TODO: handle more complex types
+            _ => None,
+        }
+    }
 }
 
 impl<'db> From<&Type<'db>> for Type<'db> {
@@ -1204,6 +1243,7 @@ mod tests {
         Unknown,
         Any,
         IntLiteral(i64),
+        BoolLiteral(bool),
         StringLiteral(&'static str),
         LiteralString,
         BytesLiteral(&'static str),
@@ -1222,6 +1262,7 @@ mod tests {
                 Ty::StringLiteral(s) => {
                     Type::StringLiteral(StringLiteralType::new(db, (*s).into()))
                 }
+                Ty::BoolLiteral(b) => Type::BooleanLiteral(b),
                 Ty::LiteralString => Type::LiteralString,
                 Ty::BytesLiteral(s) => {
                     Type::BytesLiteral(BytesLiteralType::new(db, s.as_bytes().into()))
@@ -1331,4 +1372,32 @@ mod tests {
         let db = setup_db();
         assert_eq!(ty.into_type(&db).bool(&db), Truthiness::Ambiguous);
     }
+
+    #[test_case(Ty::IntLiteral(1), Some("1"))]
+    #[test_case(Ty::BoolLiteral(true), Some("True"))]
+    #[test_case(Ty::BoolLiteral(false), Some("False"))]
+    #[test_case(Ty::StringLiteral("hello"), Some("hello"))] // no quotes
+    #[test_case(Ty::LiteralString, None)]
+    #[test_case(Ty::BuiltinInstance("int"), None)]
+    fn has_correct_str(ty: Ty, expected: Option<&str>) {
+        let db = setup_db();
+
+        let expected = expected.map(|s| Type::StringLiteral(StringLiteralType::new(&db, s.into())));
+
+        assert_eq!(ty.into_type(&db).str(&db), expected);
+    }
+
+    #[test_case(Ty::IntLiteral(1), Some("1"))]
+    #[test_case(Ty::BoolLiteral(true), Some("True"))]
+    #[test_case(Ty::BoolLiteral(false), Some("False"))]
+    #[test_case(Ty::StringLiteral("hello"), Some("'hello'"))] // single quotes
+    #[test_case(Ty::LiteralString, None)]
+    #[test_case(Ty::BuiltinInstance("int"), None)]
+    fn has_correct_repr(ty: Ty, expected: Option<&str>) {
+        let db = setup_db();
+
+        let expected = expected.map(|s| Type::StringLiteral(StringLiteralType::new(&db, s.into())));
+
+        assert_eq!(ty.into_type(&db).repr(&db), expected);
+    }
 }

crates/red_knot_python_semantic/src/types/display.rs

@@ -54,7 +54,7 @@ impl fmt::Debug for DisplayType<'_> {
 /// Writes the string representation of a type, which is the value displayed either as
 /// `Literal[<repr>]` or `Literal[<repr1>, <repr2>]` for literal types or as `<repr>` for
 /// non literals
-struct DisplayRepresentation<'db> {
+pub(super) struct DisplayRepresentation<'db> {
     ty: Type<'db>,
     db: &'db dyn Db,
 }

crates/red_knot_python_semantic/src/types/infer.rs

@@ -1653,48 +1653,71 @@ impl<'db> TypeInferenceBuilder<'db> {
     fn infer_fstring_expression(&mut self, fstring: &ast::ExprFString) -> Type<'db> {
         let ast::ExprFString { range: _, value } = fstring;
 
+        let mut done = false;
+        let mut has_expression = false;
+        let mut concatenated = String::new();
         for part in value {
             match part {
-                ast::FStringPart::Literal(_) => {
-                    // TODO string literal type
+                ast::FStringPart::Literal(literal) => {
+                    concatenated.push_str(&literal.value);
                 }
                 ast::FStringPart::FString(fstring) => {
-                    let ast::FString {
-                        range: _,
-                        elements,
-                        flags: _,
-                    } = fstring;
-                    for element in elements {
-                        self.infer_fstring_element(element);
+                    for element in &fstring.elements {
+                        match element {
+                            ast::FStringElement::Expression(expression) => {
+                                let ast::FStringExpressionElement {
+                                    range: _,
+                                    expression,
+                                    debug_text: _,
+                                    conversion,
+                                    format_spec,
+                                } = expression;
+                                // Always infer sub-expressions, even if we've figured out the type
+                                let ty = self.infer_expression(expression);
+                                if !done {
+                                    // TODO: handle format specifiers by calling a method
+                                    // (`Type::format`?) that handles the `__format__` method.
+                                    // Conversion flags should be handled before calling
+                                    // `__format__`.
+                                    // https://docs.python.org/3/library/string.html#format-string-syntax
+                                    if !conversion.is_none() || format_spec.is_some() {
+                                        has_expression = true;
+                                        done = true;
+                                    } else {
+                                        if let Some(Type::StringLiteral(literal)) = ty.str(self.db)
+                                        {
+                                            concatenated.push_str(literal.value(self.db));
+                                        } else {
+                                            has_expression = true;
+                                            done = true;
+                                        }
+                                    }
+                                }
+                            }
+                            ast::FStringElement::Literal(literal) => {
+                                if !done {
+                                    concatenated.push_str(&literal.value);
+                                }
+                            }
+                        }
                     }
                 }
             }
+            if concatenated.len() > Self::MAX_STRING_LITERAL_SIZE {
+                done = true;
+            }
         }
 
-        // TODO str type
-        Type::Unknown
-    }
-
-    fn infer_fstring_element(&mut self, element: &ast::FStringElement) {
-        match element {
-            ast::FStringElement::Literal(_) => {
-                // TODO string literal type
-            }
-            ast::FStringElement::Expression(expr_element) => {
-                let ast::FStringExpressionElement {
-                    range: _,
-                    expression,
-                    debug_text: _,
-                    conversion: _,
-                    format_spec,
-                } = expr_element;
-                self.infer_expression(expression);
-
-                if let Some(format_spec) = format_spec {
-                    for spec_element in &format_spec.elements {
-                        self.infer_fstring_element(spec_element);
-                    }
-                }
+        if has_expression {
+            builtins_symbol_ty(self.db, "str").to_instance(self.db)
+        } else {
+            if concatenated.len() <= Self::MAX_STRING_LITERAL_SIZE {
+                Type::StringLiteral(StringLiteralType::new(
+                    self.db,
+                    concatenated.into_boxed_str(),
+                ))
+            } else {
+                Type::LiteralString
             }
         }
     }
@@ -3593,6 +3616,71 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn fstring_expression() -> anyhow::Result<()> {
+        let mut db = setup_db();
+
+        db.write_dedented(
+            "src/a.py",
+            "
+            x = 0
+            y = str()
+            z = False
+
+            a = f'hello'
+            b = f'h {x}'
+            c = 'one ' f'single ' f'literal'
+            d = 'first ' f'second({b})' f' third'
+            e = f'-{y}-'
+            f = f'-{y}-' f'--' '--'
+            g = f'{z} == {False} is {True}'
+            ",
+        )?;
+
+        assert_public_ty(&db, "src/a.py", "a", "Literal[\"hello\"]");
+        assert_public_ty(&db, "src/a.py", "b", "Literal[\"h 0\"]");
+        assert_public_ty(&db, "src/a.py", "c", "Literal[\"one single literal\"]");
+        assert_public_ty(&db, "src/a.py", "d", "Literal[\"first second(h 0) third\"]");
+        assert_public_ty(&db, "src/a.py", "e", "str");
+        assert_public_ty(&db, "src/a.py", "f", "str");
+        assert_public_ty(&db, "src/a.py", "g", "Literal[\"False == False is True\"]");
+
+        Ok(())
+    }
+
+    #[test]
+    fn fstring_expression_with_conversion_flags() -> anyhow::Result<()> {
+        let mut db = setup_db();
+
+        db.write_dedented(
+            "src/a.py",
+            "
+            string = 'hello'
+            a = f'{string!r}'
+            ",
+        )?;
+
+        assert_public_ty(&db, "src/a.py", "a", "str"); // Should be `Literal["'hello'"]`
+
+        Ok(())
+    }
+
+    #[test]
+    fn fstring_expression_with_format_specifier() -> anyhow::Result<()> {
+        let mut db = setup_db();
+
+        db.write_dedented(
+            "src/a.py",
+            "
+            a = f'{1:02}'
+            ",
+        )?;
+
+        assert_public_ty(&db, "src/a.py", "a", "str"); // Should be `Literal["01"]`
+
+        Ok(())
+    }
+
     #[test]
     fn basic_call_expression() -> anyhow::Result<()> {
         let mut db = setup_db();