From f2c6cf36cf422f8ada0a02b5713d9c43c3f16304 Mon Sep 17 00:00:00 2001
From: Carl Meyer <carl@astral.sh>
Date: Thu, 30 May 2024 17:23:15 -0600
Subject: [PATCH 1/2] [red-knot] infer int literal types

---
 crates/red_knot/src/types.rs       |  6 +++++
 crates/red_knot/src/types/infer.rs | 39 ++++++++++++++++++++++++++++++
 2 files changed, 45 insertions(+)

diff --git a/crates/red_knot/src/types.rs b/crates/red_knot/src/types.rs
index 478a35f1c1658..f8b0201435555 100644
--- a/crates/red_knot/src/types.rs
+++ b/crates/red_knot/src/types.rs
@@ -36,6 +36,7 @@ pub enum Type {
     Instance(ClassTypeId),
     Union(UnionTypeId),
     Intersection(IntersectionTypeId),
+    IntLiteral(i64),
     // TODO protocols, callable types, overloads, generics, type vars
 }
 
@@ -78,6 +79,10 @@ impl Type {
                 // TODO return the intersection of those results
                 todo!("attribute lookup on Intersection type")
             }
+            Type::IntLiteral(_) => {
+                // TODO raise error
+                Ok(Some(Type::Unknown))
+            }
         }
     }
 }
@@ -616,6 +621,7 @@ impl std::fmt::Display for DisplayType<'_> {
                 .get_module(int_id.file_id)
                 .get_intersection(int_id.intersection_id)
                 .display(f, self.store),
+            Type::IntLiteral(n) => write!(f, "Literal[{n}]"),
         }
     }
 }
diff --git a/crates/red_knot/src/types/infer.rs b/crates/red_knot/src/types/infer.rs
index 0d6d23b8ce779..7b930fb0894b3 100644
--- a/crates/red_knot/src/types/infer.rs
+++ b/crates/red_knot/src/types/infer.rs
@@ -145,6 +145,15 @@ fn infer_expr_type(db: &dyn SemanticDb, file_id: FileId, expr: &ast::Expr) -> Qu
     // TODO cache the resolution of the type on the node
     let symbols = symbol_table(db, file_id)?;
     match expr {
+        ast::Expr::NumberLiteral(ast::ExprNumberLiteral { value: v, .. }) => {
+            if let ast::Number::Int(n) = v {
+                // TODO support big int literals, or at least default to `builtins.int`
+                Ok(n.as_i64().map(Type::IntLiteral).unwrap_or(Type::Unknown))
+            } else {
+                // TODO builtins.float or builtins.complex
+                Ok(Type::Unknown)
+            }
+        }
         ast::Expr::Name(name) => {
             // TODO look up in the correct scope, don't assume global
             if let Some(symbol_id) = symbols.root_symbol_id_by_name(&name.id) {
@@ -348,4 +357,34 @@ mod tests {
         assert_eq!(format!("{}", ty.display(&jar.type_store)), "Literal[C]");
         Ok(())
     }
+
+    #[test]
+    fn resolve_literal() -> anyhow::Result<()> {
+        let case = create_test()?;
+        let db = &case.db;
+
+        let path = case.src.path().join("a.py");
+        std::fs::write(path, "x = 1")?;
+        let file = resolve_module(db, ModuleName::new("a"))?
+            .expect("module should be found")
+            .path(db)?
+            .file();
+        let syms = symbol_table(db, file)?;
+        let x_sym = syms
+            .root_symbol_id_by_name("x")
+            .expect("x symbol should be found");
+
+        let ty = infer_symbol_type(
+            db,
+            GlobalSymbolId {
+                file_id: file,
+                symbol_id: x_sym,
+            },
+        )?;
+
+        let jar = HasJar::<SemanticJar>::jar(db)?;
+        assert!(matches!(ty, Type::IntLiteral(_)));
+        assert_eq!(format!("{}", ty.display(&jar.type_store)), "Literal[1]");
+        Ok(())
+    }
 }

From 1bed66ac17cd53383cded19146d7ab39fc972193 Mon Sep 17 00:00:00 2001
From: Carl Meyer <carl@astral.sh>
Date: Fri, 31 May 2024 13:38:43 -0600
Subject: [PATCH 2/2] use match instead of if let

---
 crates/red_knot/src/types/infer.rs | 13 +++++++------
 1 file changed, 7 insertions(+), 6 deletions(-)

diff --git a/crates/red_knot/src/types/infer.rs b/crates/red_knot/src/types/infer.rs
index 7b930fb0894b3..ff27f25c2dd09 100644
--- a/crates/red_knot/src/types/infer.rs
+++ b/crates/red_knot/src/types/infer.rs
@@ -145,13 +145,14 @@ fn infer_expr_type(db: &dyn SemanticDb, file_id: FileId, expr: &ast::Expr) -> Qu
     // TODO cache the resolution of the type on the node
     let symbols = symbol_table(db, file_id)?;
     match expr {
-        ast::Expr::NumberLiteral(ast::ExprNumberLiteral { value: v, .. }) => {
-            if let ast::Number::Int(n) = v {
-                // TODO support big int literals, or at least default to `builtins.int`
-                Ok(n.as_i64().map(Type::IntLiteral).unwrap_or(Type::Unknown))
-            } else {
+        ast::Expr::NumberLiteral(ast::ExprNumberLiteral { value, .. }) => {
+            match value {
+                ast::Number::Int(n) => {
+                    // TODO support big int literals
+                    Ok(n.as_i64().map(Type::IntLiteral).unwrap_or(Type::Unknown))
+                }
                 // TODO builtins.float or builtins.complex
-                Ok(Type::Unknown)
+                _ => Ok(Type::Unknown),
             }
         }
         ast::Expr::Name(name) => {