[red-knot] infer int literal types (#11623)

carljm · web-flow · commit 16a926d13810 · 2024-05-31T13:52:29.000-06:00
## Summary

Give red-knot the ability to infer int literal types. This is quick and
easy, mostly because these types are a convenient way to observe
control-flow handling with simple assignments.

## Test Plan

Added test.
diff --git 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
@@ -145,6 +145,16 @@ 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, .. }) => {
+            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),
+            }
+        }
         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 +358,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(())
+    }
 }

Original file line number	Diff line number	Diff line change
`@@ -36,6 +36,7 @@ pub enum Type {`
`36`	`36`	`Instance(ClassTypeId),`
`37`	`37`	`Union(UnionTypeId),`
`38`	`38`	`Intersection(IntersectionTypeId),`
	`39`	`+ IntLiteral(i64),`
`39`	`40`	`// TODO protocols, callable types, overloads, generics, type vars`
`40`	`41`	`}`
`41`	`42`
`@@ -78,6 +79,10 @@ impl Type {`
`78`	`79`	`// TODO return the intersection of those results`
`79`	`80`	`todo!("attribute lookup on Intersection type")`
`80`	`81`	`}`
	`82`	`+ Type::IntLiteral(_) => {`
	`83`	`+ // TODO raise error`
	`84`	`+ Ok(Some(Type::Unknown))`
	`85`	`+ }`
`81`	`86`	`}`
`82`	`87`	`}`
`83`	`88`	`}`
`@@ -616,6 +621,7 @@ impl std::fmt::Display for DisplayType<'_> {`
`616`	`621`	`.get_module(int_id.file_id)`
`617`	`622`	`.get_intersection(int_id.intersection_id)`
`618`	`623`	`.display(f, self.store),`
	`624`	`+ Type::IntLiteral(n) => write!(f, "Literal[{n}]"),`
`619`	`625`	`}`
`620`	`626`	`}`
`621`	`627`	`}`