Handle inferring Literal type annotation

crates/red_knot_python_semantic/resources/mdtest/comparison/instances/membership_test.md

@@ -95,13 +95,11 @@ class AlwaysFalse:
     def __contains__(self, item: int) -> Literal[""]:
         return ""
 
-# TODO: it should be Literal[True] and Literal[False]
-reveal_type(42 in AlwaysTrue())  # revealed: @Todo
-reveal_type(42 not in AlwaysTrue())  # revealed: @Todo
+reveal_type(42 in AlwaysTrue())  # revealed: Literal[True]
+reveal_type(42 not in AlwaysTrue())  # revealed: Literal[False]
 
-# TODO: it should be Literal[False] and Literal[True]
-reveal_type(42 in AlwaysFalse())  # revealed: @Todo
-reveal_type(42 not in AlwaysFalse())  # revealed: @Todo
+reveal_type(42 in AlwaysFalse())  # revealed: Literal[False]
+reveal_type(42 not in AlwaysFalse())  # revealed: Literal[True]
 ```
 
 ## No Fallback for `__contains__`

crates/red_knot_python_semantic/resources/mdtest/literal/literal.md

@@ -3,24 +3,37 @@
 ```py
 from typing import Literal
 
-x: Literal[10] = 10
 mode: Literal["w", "r"] = "w"
-mode2: Literal["w"] | Literal["r"] = "w"
 
-# reveal_type(Literal[26])  # revealed: Literal[26]
-# reveal_type(Literal[0x1A])  # revealed: Literal[26]
-# reveal_type(Literal[-4])  # revealed: Literal[-4]
-# reveal_type(Literal["hello world"])  # revealed: Literal["hello world"]
-# reveal_type(Literal[b"hello world"])  # revealed: Literal[b"hello world"]
-# reveal_type(Literal["hello world"])  # revealed: Literal["hello world"]
-# reveal_type(Literal[True])  # revealed: Literal[True]
-# reveal_type(Literal[Color.RED]) # revealed: Literal["Red"]
-# reveal_type(Literal[None])  # revealed: None
-# TODO: "Revealed type is `Literal[1, 2, 3] | Literal["foo"] | Literal[5] | None`"
-# reveal_type(Literal[Literal[Literal[1, 2, 3], "foo"], 5, None]) # revealed: Literal[1, 2, 3, "foo", 5, None]
+# TODO: PEP-604 unions should not give error
+# error: [unsupported-operator] "Operator `|` is unsupported between objects of type `object` and `object`"
+mode2: Literal["w"] | Literal["r"] = "w"
+# union_var: Literal[Literal[Literal[1, 2, 3], "foo"], 5, None]
+# reveal_type(union_var) # revealed: Literal[1, 2, 3, "foo", 5, None]
 
 def f():
-    reveal_type(x)  # revealed: Literal[10]
     reveal_type(mode)  # revealed: Literal["w", "r"]
-    # reveal_type(mode2)  # revealed: Literal["w", "r"]
+    reveal_type(mode2)  # revealed: @Todo
+
+a: Literal[26] = 26
+reveal_type(a)  # revealed: Literal[26]
+a2: Literal[0x1A] = 0x1A
+reveal_type(a2)  # revealed: Literal[26]
+a3: Literal[-4] = -4
+reveal_type(a3)  # revealed: Literal[-4]
+a4: Literal["hello world"] = "hello world"
+reveal_type(a4)  # revealed: Literal["hello world"]
+a5: Literal[b"hello world"] = b"hello world"
+reveal_type(a5)  # revealed: Literal[b"hello world"]
+a6: Literal["hello world"] = "hello world"
+reveal_type(a6)  # revealed: Literal["hello world"]
+a7: Literal[True] = True
+reveal_type(a7)  # revealed: Literal[True]
+# a: Literal[Color.RED]
+# reveal_type(a)
+a8: Literal[None] = None
+reveal_type(a8)  # revealed: None
+
+# error: [invalid-literal-parameter] "Type arguments for `Literal` must be None, a literal value (int, bool, str, or bytes), or an enum value"
+a9: Literal[3 + 4] = 7
 ```

crates/red_knot_python_semantic/resources/mdtest/unary/instance.md

@@ -1,6 +1,8 @@
 # Unary Operations
 
 ```py
+from typing import Literal
+
 class Number:
     def __init__(self, value: int):
         self.value = 1
@@ -19,7 +21,7 @@ a = Number()
 reveal_type(+a)  # revealed: int
 reveal_type(-a)  # revealed: int
 # TODO: this should be True
-reveal_type(~a)  # revealed: object
+reveal_type(~a)  # revealed: Literal[True]
 
 class NoDunder: ...
 

crates/red_knot_python_semantic/src/types.rs

@@ -1047,10 +1047,16 @@ impl<'db> Type<'db> {
             Type::Intersection(_) => Type::Todo,
             // TODO: calling `.to_instance()` on any of these should result in a diagnostic,
             // since they already indicate that the object is an instance of some kind:
+            Type::Instance(instance) => {
+                if instance.known_instance(db).is_some() {
+                    *self
+                } else {
+                    Type::Unknown
+                }
+            }
             Type::BooleanLiteral(_)
             | Type::BytesLiteral(_)
             | Type::FunctionLiteral(_)
-            | Type::Instance(_)
             | Type::ModuleLiteral(_)
             | Type::IntLiteral(_)
             | Type::StringLiteral(_)
@@ -1279,6 +1285,14 @@ pub enum KnownInstance {
     Literal,
 }
 
+impl<'db> KnownInstance {
+    pub const fn as_str(&self) -> &'static str {
+        match self {
+            KnownInstance::Literal => "Literal",
+        }
+    }
+}
+
 #[derive(Debug, Clone, PartialEq, Eq)]
 enum CallOutcome<'db> {
     Callable {

crates/red_knot_python_semantic/src/types/infer.rs

@@ -1380,19 +1380,13 @@ impl<'db> TypeInferenceBuilder<'db> {
             let class = instance.class_type(self.db);
             if class.is_known(self.db, KnownClass::SpecialForm) {
                 if let Some(name_expr) = target.as_name_expr() {
-                    let target_name = name_expr.id.clone();
-                    let new_class = ClassType::new(
-                        self.db,
-                        target_name,
-                        class.definition(self.db),
-                        class.body_scope(self.db),
-                        Some(KnownClass::SpecialForm),
-                    );
-                    annotation_ty = Type::Instance(InstanceType::new(
-                        self.db,
-                        new_class,
-                        Some(KnownInstance::Literal),
-                    ));
+                    if name_expr.id == KnownInstance::Literal.as_str() {
+                        annotation_ty = Type::Instance(InstanceType::new(
+                            self.db,
+                            class,
+                            Some(KnownInstance::Literal),
+                        ));
+                    }
                 }
             }
         }
@@ -3477,29 +3471,60 @@ impl<'db> TypeInferenceBuilder<'db> {
         } = subscript;
 
         let value_ty = self.infer_type_expression(value);
-        let slice_ty = self.infer_type_expression(slice);
 
         // Handling of Special Forms
-        match (value_ty, slice_ty) {
-            (Type::Instance(class), slice_ty)
-                if class.is_known(self.db, KnownClass::SpecialForm) =>
-            {
-                match class.known_instance(self.db) {
-                    Some(s) => match s {
-                        KnownInstance::Literal => match slice_ty {
-                            // TODO: not correct logic
+        match value_ty {
+            Type::Instance(class) if class.is_known(self.db, KnownClass::SpecialForm) => {
+                let Some(s) = class.known_instance(self.db) else {
+                    return Type::Todo;
+                };
+                // NOTE: slice_ty is treated as expression because Literal accepts expression
+                // inside the []
+                let slice_ty = self.infer_expression(slice);
+                match s {
+                    KnownInstance::Literal => {
+                        match **slice {
+                    ruff_python_ast::Expr::StringLiteral(_)
+                    | ruff_python_ast::Expr::BytesLiteral(_)
+                    | ruff_python_ast::Expr::NumberLiteral(_)
+                    | ruff_python_ast::Expr::BooleanLiteral(_)
+                    | ruff_python_ast::Expr::Tuple(_)
+                    // For enum values
+                    | ruff_python_ast::Expr::Attribute(_)
+                    // For Another Literal inside this Literal
+                    | ruff_python_ast::Expr::Subscript(_)
+                    | ruff_python_ast::Expr::NoneLiteral(_) => {}
+                    // for negative numbers
+                    ruff_python_ast::Expr::UnaryOp(ref u) if (u.op == UnaryOp::USub || u.op == UnaryOp::UAdd) && u.operand.is_number_literal_expr() => {}
+                    _ => {
+                        self.add_diagnostic(
+                        slice.as_ref().into(),
+                        "invalid-literal-parameter",
+                        format_args!(
+                            "Type arguments for `Literal` must be None, a literal value (int, bool, str, or bytes), or an enum value",
+                        ),
+                    );
+                        return Type::Unknown;
+                    }
+                    };
+                        match slice_ty {
                             Type::Tuple(tuple) => {
                                 let elts = tuple.elements(self.db);
                                 Type::Union(UnionType::new(self.db, elts))
                             }
                             ty => ty,
-                        },
-                    },
-                    None => todo!(),
+                        }
+                    }
                 }
             }
-            // TODO: emit diagnostic
-            _ => Type::Todo,
+            value_ty => {
+                let value_node = value.as_ref();
+                let slice_ty = self.infer_expression(slice);
+                // TODO: currently the logic to get the type of type of a subscript in type
+                // annotation with where value is a slice is defined in the
+                // subscript_expression_types. Once it's complete call that logic from here.
+                self.infer_subscript_expression_types(value_node, value_ty, slice_ty)
+            }
         }
     }
 }
@@ -3528,8 +3553,6 @@ impl<'db> TypeInferenceBuilder<'db> {
 impl<'db> TypeInferenceBuilder<'db> {
     fn infer_type_expression(&mut self, expression: &ast::Expr) -> Type<'db> {
         // https://typing.readthedocs.io/en/latest/spec/annotations.html#grammar-token-expression-grammar-type_expression
-        // TODO: this does not include any of the special forms, and is only a
-        //   stub of the forms other than a standalone name in scope.
 
         let ty = match expression {
             ast::Expr::Name(name) => {
@@ -3835,13 +3858,17 @@ fn perform_rich_comparison<'db>(
     // TODO: this currently gives the return type even if the arg types are invalid
     // (e.g. int.__lt__ with string instance should be errored, currently bool)
 
-    let left_instance = Type::Instance(InstanceType::new(db, left_class, None));
-    let right_instance = Type::Instance(InstanceType::new(db, right_class, None));
     let call_dunder =
         |op: RichCompareOperator, left_class: ClassType<'db>, right_class: ClassType<'db>| {
             left_class
                 .class_member(db, op.dunder())
-                .call(db, &[left_instance, right_instance])
+                .call(
+                    db,
+                    &[
+                        Type::Instance(InstanceType::new(db, left_class, None)),
+                        Type::Instance(InstanceType::new(db, right_class, None)),
+                    ],
+                )
                 .return_ty(db)
         };
 
@@ -3865,8 +3892,8 @@ fn perform_rich_comparison<'db>(
     })
     .ok_or_else(|| CompareUnsupportedError {
         op: op.into(),
-        left_ty: left_instance,
-        right_ty: right_instance,
+        left_ty: Type::Instance(InstanceType::new(db, left_class, None)),
+        right_ty: Type::Instance(InstanceType::new(db, right_class, None)),
     })
 }