[red-knot] Deeper understanding of LiteralString

crates/red_knot_python_semantic/resources/mdtest/annotations/literal_string.md

@@ -0,0 +1,128 @@
+# `LiteralString`
+
+`LiteralString` represents a string that is either defined directly within the source code or is
+made up of such components.
+
+Parts of the testcases defined here were adapted from [the specification's examples][1].
+
+## Usages
+
+### Valid places
+
+It can be used anywhere a type is accepted:
+
+```py
+from typing import LiteralString
+
+x: LiteralString
+
+def f():
+    reveal_type(x)  # revealed: LiteralString
+```
+
+### Within `Literal`
+
+`LiteralString` cannot be used within `Literal`:
+
+```py
+from typing import Literal, LiteralString
+
+bad_union: Literal["hello", LiteralString]  # error: [invalid-literal-parameter]
+bad_nesting: Literal[LiteralString]  # error: [invalid-literal-parameter]
+```
+
+### Parametrized
+
+`LiteralString` cannot be parametrized.
+
+```py
+from typing import LiteralString
+
+a: LiteralString[str]  # error: [invalid-type-parameter]
+b: LiteralString["foo"]  # error: [invalid-type-parameter]
+```
+
+### As a base class
+
+Subclassing `LiteralString` leads to a runtime error.
+
+```py
+from typing import LiteralString
+
+class C(LiteralString): ...  # error: [invalid-base]
+```
+
+## Inference
+
+### Common operations
+
+```py
+foo: LiteralString = "foo"
+reveal_type(foo)  # revealed: Literal["foo"]
+
+bar: LiteralString = "bar"
+reveal_type(foo + bar)  # revealed: Literal["foobar"]
+
+baz: LiteralString = "baz"
+baz += foo
+reveal_type(baz)  # revealed: Literal["bazfoo"]
+
+qux = (foo, bar)
+reveal_type(qux)  # revealed: tuple[Literal["foo"], Literal["bar"]]
+
+# TODO: Infer "LiteralString"
+reveal_type(foo.join(qux))  # revealed: @Todo(call todo)
+
+template: LiteralString = "{}, {}"
+reveal_type(template)  # revealed: Literal["{}, {}"]
+# TODO: Infer `LiteralString`
+reveal_type(template.format(foo, bar))  # revealed: @Todo(call todo)
+```
+
+### Assignability
+
+`Literal[""]` is assignable to `LiteralString`, and `LiteralString` is assignable to `str`, but not
+vice versa.
+
+```py
+def coinflip() -> bool:
+    return True
+
+foo_1: Literal["foo"] = "foo"
+bar_1: LiteralString = foo_1  # fine
+
+foo_2 = "foo" if coinflip() else "bar"
+reveal_type(foo_2)  # revealed: Literal["foo", "bar"]
+bar_2: LiteralString = foo_2  # fine
+
+foo_3: LiteralString = "foo" * 1_000_000_000
+bar_3: str = foo_2  # fine
+
+baz_1: str = str()
+qux_1: LiteralString = baz_1  # error: [invalid-assignment]
+
+baz_2: LiteralString = "baz" * 1_000_000_000
+qux_2: Literal["qux"] = baz_2  # error: [invalid-assignment]
+
+baz_3 = "foo" if coinflip() else 1
+reveal_type(baz_3)  # revealed: Literal["foo"] | Literal[1]
+qux_3: LiteralString = baz_3  # error: [invalid-assignment]
+```
+
+### Narrowing
+
+```py
+lorem: LiteralString = "lorem" * 1_000_000_000
+
+reveal_type(lorem)  # revealed: LiteralString
+
+if lorem == "ipsum":
+    reveal_type(lorem)  # revealed: Literal["ipsum"]
+
+reveal_type(lorem)  # revealed: LiteralString
+
+if "" < lorem == "ipsum":
+    reveal_type(lorem)  # revealed: Literal["ipsum"]
+```
+
+[1]: https://typing.readthedocs.io/en/latest/spec/literal.html#literalstring

crates/red_knot_python_semantic/src/types.rs

@@ -1403,7 +1403,7 @@ impl<'db> Type<'db> {
             // `Any` is callable, and its return type is also `Any`.
             Type::Any => CallOutcome::callable(Type::Any),
 
-            Type::Todo(_) => CallOutcome::callable(todo_type!()),
+            Type::Todo(_) => CallOutcome::callable(todo_type!("call todo")),
 
             Type::Unknown => CallOutcome::callable(Type::Unknown),
 
@@ -1556,6 +1556,7 @@ impl<'db> Type<'db> {
             Type::KnownInstance(KnownInstanceType::Never | KnownInstanceType::NoReturn) => {
                 Type::Never
             }
+            Type::KnownInstance(KnownInstanceType::LiteralString) => Type::LiteralString,
             _ => todo_type!(),
         }
     }
@@ -1888,6 +1889,8 @@ impl<'db> KnownClass {
 pub enum KnownInstanceType<'db> {
     /// The symbol `typing.Literal` (which can also be found as `typing_extensions.Literal`)
     Literal,
+    /// The symbol `typing.LiteralString` (which can also be found as `typing_extensions.LiteralString`)
+    LiteralString,
     /// The symbol `typing.Optional` (which can also be found as `typing_extensions.Optional`)
     Optional,
     /// The symbol `typing.Union` (which can also be found as `typing_extensions.Union`)
@@ -1907,6 +1910,7 @@ impl<'db> KnownInstanceType<'db> {
     pub const fn as_str(self) -> &'static str {
         match self {
             Self::Literal => "Literal",
+            Self::LiteralString => "LiteralString",
             Self::Optional => "Optional",
             Self::Union => "Union",
             Self::TypeVar(_) => "TypeVar",
@@ -1920,6 +1924,7 @@ impl<'db> KnownInstanceType<'db> {
     pub const fn bool(self) -> Truthiness {
         match self {
             Self::Literal
+            | Self::LiteralString
             | Self::Optional
             | Self::TypeVar(_)
             | Self::Union
@@ -1933,6 +1938,7 @@ impl<'db> KnownInstanceType<'db> {
     pub fn repr(self, db: &'db dyn Db) -> &'db str {
         match self {
             Self::Literal => "typing.Literal",
+            Self::LiteralString => "typing.LiteralString",
             Self::Optional => "typing.Optional",
             Self::Union => "typing.Union",
             Self::NoReturn => "typing.NoReturn",
@@ -1946,6 +1952,7 @@ impl<'db> KnownInstanceType<'db> {
     pub const fn class(self) -> KnownClass {
         match self {
             Self::Literal => KnownClass::SpecialForm,
+            Self::LiteralString => KnownClass::SpecialForm,
             Self::Optional => KnownClass::SpecialForm,
             Self::Union => KnownClass::SpecialForm,
             Self::NoReturn => KnownClass::SpecialForm,
@@ -1970,6 +1977,7 @@ impl<'db> KnownInstanceType<'db> {
         }
         match (module.name().as_str(), instance_name) {
             ("typing" | "typing_extensions", "Literal") => Some(Self::Literal),
+            ("typing" | "typing_extensions", "LiteralString") => Some(Self::LiteralString),
             ("typing" | "typing_extensions", "Optional") => Some(Self::Optional),
             ("typing" | "typing_extensions", "Union") => Some(Self::Union),
             ("typing" | "typing_extensions", "NoReturn") => Some(Self::NoReturn),

crates/red_knot_python_semantic/src/types/infer.rs

@@ -4642,6 +4642,17 @@ impl<'db> TypeInferenceBuilder<'db> {
                 );
                 Type::Unknown
             }
+            KnownInstanceType::LiteralString => {
+                self.diagnostics.add(
+                    subscript.into(),
+                    "invalid-type-parameter",
+                    format_args!(
+                        "Type `{}` expected no type parameter. Did you mean to use `Literal[...]` instead?",
+                        known_instance.repr(self.db)
+                    ),
+                );
+                Type::Unknown
+            }
         }
     }
 

crates/red_knot_python_semantic/src/types/mro.rs

@@ -374,6 +374,7 @@ impl<'db> ClassBase<'db> {
                 KnownInstanceType::TypeVar(_)
                 | KnownInstanceType::TypeAliasType(_)
                 | KnownInstanceType::Literal
+                | KnownInstanceType::LiteralString
                 | KnownInstanceType::Union
                 | KnownInstanceType::NoReturn
                 | KnownInstanceType::Never

crates/red_knot_python_semantic/src/types/narrow.rs

@@ -291,8 +291,15 @@ impl<'db> NarrowingConstraintsBuilder<'db> {
             .chain(comparators)
             .tuple_windows::<(&ruff_python_ast::Expr, &ruff_python_ast::Expr)>();
         let mut constraints = NarrowingConstraints::default();
+
+        let mut last_rhs_ty: Option<Type> = None;
+
         for (op, (left, right)) in std::iter::zip(&**ops, comparator_tuples) {
+            let lhs_ty = last_rhs_ty.unwrap_or_else(|| {
+                inference.expression_ty(left.scoped_expression_id(self.db, scope))
+            });
             let rhs_ty = inference.expression_ty(right.scoped_expression_id(self.db, scope));
+            last_rhs_ty = Some(rhs_ty);
 
             match left {
                 ast::Expr::Name(ast::ExprName {
@@ -327,6 +334,9 @@ impl<'db> NarrowingConstraintsBuilder<'db> {
                                 constraints.insert(symbol, ty);
                             }
                         }
+                        ast::CmpOp::Eq if lhs_ty.is_literal_string() => {
+                            constraints.insert(symbol, rhs_ty);
+                        }
                         _ => {
                             // TODO other comparison types
                         }