Handle Literal special form

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

@@ -0,0 +1,27 @@
+## Literal
+
+```py
+from typing import Literal
+
+x: Literal[10] = 10
+mode: Literal["w", "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]
+# TODO: check enum
+# 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]
+
+
+def f():
+    # currently we reveal `@Todo` here
+    reveal_type(x)  # revealed: Literal[10]
+    reveal_type(mode)  # revealed: Literal["w", "r"]
+```

crates/red_knot_python_semantic/src/types.rs

@@ -1110,6 +1110,7 @@ pub enum KnownClass {
     FunctionType,
     // Typeshed
     NoneType, // Part of `types` for Python >= 3.10
+    SpecialForm,
 }
 
 impl<'db> KnownClass {
@@ -1130,6 +1131,7 @@ impl<'db> KnownClass {
             Self::ModuleType => "ModuleType",
             Self::FunctionType => "FunctionType",
             Self::NoneType => "NoneType",
+            Self::SpecialForm => "SpecialForm",
         }
     }
 
@@ -1154,6 +1156,11 @@ impl<'db> KnownClass {
                 types_symbol_ty(db, self.as_str())
             }
             Self::NoneType => typeshed_symbol_ty(db, self.as_str()),
+            Self::SpecialForm => {
+                let t = typeshed_symbol_ty(db, self.as_str());
+                debug_assert!(!t.is_unbound(), "special form not found");
+                t
+            }
         }
     }
 
@@ -1186,6 +1193,7 @@ impl<'db> KnownClass {
             "NoneType" => Some(Self::NoneType),
             "ModuleType" => Some(Self::ModuleType),
             "FunctionType" => Some(Self::FunctionType),
+            "_SpecialForm" => Some(Self::SpecialForm),
             _ => None,
         }
     }
@@ -1209,6 +1217,9 @@ impl<'db> KnownClass {
             | Self::Dict => module.name() == "builtins",
             Self::GenericAlias | Self::ModuleType | Self::FunctionType => module.name() == "types",
             Self::NoneType => matches!(module.name().as_str(), "_typeshed" | "types"),
+            Self::SpecialForm => {
+                matches!(module.name().as_str(), "typing")
+            }
         }
     }
 }

crates/red_knot_python_semantic/src/types/infer.rs

@@ -1364,7 +1364,25 @@ impl<'db> TypeInferenceBuilder<'db> {
             simple: _,
         } = assignment;
 
-        let annotation_ty = self.infer_annotation_expression(annotation);
+        let mut annotation_ty = self.infer_annotation_expression(annotation);
+
+        // If the variable is annotation with SpecialForm then create a new class with name of the
+        // variable.
+        if let Type::Instance(class) = annotation_ty {
+            if class.is_known(self.db, KnownClass::SpecialForm) {
+                let target_name = target.as_name_expr().unwrap().id.clone();
+                let new_class = ClassType::new(
+                    self.db,
+                    target_name,
+                    class.definition(self.db),
+                    class.body_scope(self.db),
+                    // We say this is a known class and is a special form
+                    Some(KnownClass::SpecialForm),
+                );
+                annotation_ty = Type::Class(new_class);
+            }
+        }
+
         if let Some(value) = value {
             let value_ty = self.infer_expression(value);
             self.add_declaration_with_binding(
@@ -3276,6 +3294,21 @@ impl<'db> TypeInferenceBuilder<'db> {
                     value_ty,
                     Type::IntLiteral(i64::from(bool)),
                 ),
+
+            // Handling of Special Forms
+            (Type::Class(class), slice_ty) if class.is_known(self.db, KnownClass::SpecialForm) => {
+                if class.name(self.db) == "Literal" {
+                    match slice_ty {
+                        Type::Tuple(tuple) => {
+                            let elts = tuple.elements(self.db);
+                            Type::Union(UnionType::new(self.db, elts))
+                        }
+                        ty => ty,
+                    }
+                } else {
+                    Type::Todo
+                }
+            }
             (value_ty, slice_ty) => {
                 // Resolve the value to its class.
                 let value_meta_ty = value_ty.to_meta_type(self.db);
@@ -3492,11 +3525,7 @@ impl<'db> TypeInferenceBuilder<'db> {
             ast::Expr::NumberLiteral(_literal) => Type::Todo,
             ast::Expr::BooleanLiteral(_literal) => Type::Todo,
 
-            // TODO: this may be a place we need to revisit with special forms.
-            ast::Expr::Subscript(subscript) => {
-                self.infer_subscript_expression(subscript);
-                Type::Todo
-            }
+            ast::Expr::Subscript(subscript) => self.infer_subscript_expression(subscript),
 
             // Forms which are invalid in the context of annotation expressions: we infer their
             // nested expressions as normal expressions, but the type of the top-level expression is