[red-knot] infer int literal types

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}]"),
         }
     }
 }

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(())
+    }
 }