[red-knot] Resolve symbols from builtins.pyi in the stdlib if they cannot be found in other scopes (#12390)

Co-authored-by: Carl Meyer <carl@astral.sh>

crates/red_knot_python_semantic/src/builtins.rs

@@ -0,0 +1,16 @@
+use red_knot_module_resolver::{resolve_module, ModuleName};
+
+use crate::semantic_index::global_scope;
+use crate::semantic_index::symbol::ScopeId;
+use crate::Db;
+
+/// Salsa query to get the builtins scope.
+///
+/// Can return None if a custom typeshed is used that is missing `builtins.pyi`.
+#[salsa::tracked]
+pub(crate) fn builtins_scope(db: &dyn Db) -> Option<ScopeId<'_>> {
+    let builtins_name =
+        ModuleName::new_static("builtins").expect("Expected 'builtins' to be a valid module name");
+    let builtins_file = resolve_module(db.upcast(), builtins_name)?.file();
+    Some(global_scope(db, builtins_file))
+}

crates/red_knot_python_semantic/src/db.rs

@@ -3,6 +3,7 @@ use salsa::DbWithJar;
 use red_knot_module_resolver::Db as ResolverDb;
 use ruff_db::{Db as SourceDb, Upcast};
 
+use crate::builtins::builtins_scope;
 use crate::semantic_index::definition::Definition;
 use crate::semantic_index::expression::Expression;
 use crate::semantic_index::symbol::ScopeId;
@@ -28,6 +29,7 @@ pub struct Jar(
     infer_definition_types,
     infer_expression_types,
     infer_scope_types,
+    builtins_scope,
 );
 
 /// Database giving access to semantic information about a Python program.

crates/red_knot_python_semantic/src/lib.rs

@@ -6,6 +6,7 @@ pub use db::{Db, Jar};
 pub use semantic_model::{HasTy, SemanticModel};
 
 pub mod ast_node_ref;
+mod builtins;
 mod db;
 mod node_key;
 pub mod semantic_index;

crates/red_knot_python_semantic/src/types.rs

@@ -1,6 +1,7 @@
 use ruff_db::files::File;
 use ruff_python_ast::name::Name;
 
+use crate::builtins::builtins_scope;
 use crate::semantic_index::definition::Definition;
 use crate::semantic_index::symbol::{ScopeId, ScopedSymbolId};
 use crate::semantic_index::{global_scope, symbol_table, use_def_map};
@@ -47,6 +48,15 @@ pub(crate) fn global_symbol_ty_by_name<'db>(db: &'db dyn Db, file: File, name: &
     symbol_ty_by_name(db, global_scope(db, file), name)
 }
 
+/// Shorthand for `symbol_ty` that looks up a symbol in the builtins.
+///
+/// Returns `None` if the builtins module isn't available for some reason.
+pub(crate) fn builtins_symbol_ty_by_name<'db>(db: &'db dyn Db, name: &str) -> Type<'db> {
+    builtins_scope(db)
+        .map(|builtins| symbol_ty_by_name(db, builtins, name))
+        .unwrap_or(Type::Unbound)
+}
+
 /// Infer the type of a [`Definition`].
 pub(crate) fn definition_ty<'db>(db: &'db dyn Db, definition: Definition<'db>) -> Type<'db> {
     let inference = infer_definition_types(db, definition);

crates/red_knot_python_semantic/src/types/display.rs

@@ -29,13 +29,9 @@ impl Display for DisplayType<'_> {
                 write!(f, "<module '{:?}'>", file.path(self.db.upcast()))
             }
             // TODO functions and classes should display using a fully qualified name
-            Type::Class(class) => {
-                f.write_str("Literal[")?;
-                f.write_str(&class.name(self.db))?;
-                f.write_str("]")
-            }
+            Type::Class(class) => write!(f, "Literal[{}]", class.name(self.db)),
             Type::Instance(class) => f.write_str(&class.name(self.db)),
-            Type::Function(function) => f.write_str(&function.name(self.db)),
+            Type::Function(function) => write!(f, "Literal[{}]", function.name(self.db)),
             Type::Union(union) => union.display(self.db).fmt(f),
             Type::Intersection(intersection) => intersection.display(self.db).fmt(f),
             Type::IntLiteral(n) => write!(f, "Literal[{n}]"),

crates/red_knot_python_semantic/src/types/infer.rs

@@ -29,15 +29,16 @@ use ruff_db::parsed::parsed_module;
 use ruff_python_ast as ast;
 use ruff_python_ast::{ExprContext, TypeParams};
 
+use crate::builtins::builtins_scope;
 use crate::semantic_index::ast_ids::{HasScopedAstId, HasScopedUseId, ScopedExpressionId};
 use crate::semantic_index::definition::{Definition, DefinitionKind, DefinitionNodeKey};
 use crate::semantic_index::expression::Expression;
 use crate::semantic_index::semantic_index;
-use crate::semantic_index::symbol::NodeWithScopeKind;
-use crate::semantic_index::symbol::{NodeWithScopeRef, ScopeId};
+use crate::semantic_index::symbol::{FileScopeId, NodeWithScopeKind, NodeWithScopeRef, ScopeId};
 use crate::semantic_index::SemanticIndex;
 use crate::types::{
-    definitions_ty, global_symbol_ty_by_name, ClassType, FunctionType, Name, Type, UnionTypeBuilder,
+    builtins_symbol_ty_by_name, definitions_ty, global_symbol_ty_by_name, ClassType, FunctionType,
+    Name, Type, UnionTypeBuilder,
 };
 use crate::Db;
 
@@ -686,7 +687,18 @@ impl<'db> TypeInferenceBuilder<'db> {
                     let symbol = symbols.symbol_by_name(id).unwrap();
                     if !symbol.is_defined() || !self.scope.is_function_like(self.db) {
                         // implicit global
-                        Some(global_symbol_ty_by_name(self.db, self.file, id))
+                        let mut unbound_ty = if file_scope_id == FileScopeId::global() {
+                            Type::Unbound
+                        } else {
+                            global_symbol_ty_by_name(self.db, self.file, id)
+                        };
+                        // fallback to builtins
+                        if matches!(unbound_ty, Type::Unbound)
+                            && Some(self.scope) != builtins_scope(self.db)
+                        {
+                            unbound_ty = builtins_symbol_ty_by_name(self.db, id);
+                        }
+                        Some(unbound_ty)
                     } else {
                         Some(Type::Unbound)
                     }
@@ -792,6 +804,7 @@ mod tests {
     use ruff_db::testing::assert_function_query_was_not_run;
     use ruff_python_ast::name::Name;
 
+    use crate::builtins::builtins_scope;
     use crate::db::tests::TestDb;
     use crate::semantic_index::definition::Definition;
     use crate::semantic_index::semantic_index;
@@ -819,6 +832,23 @@ mod tests {
         db
     }
 
+    fn setup_db_with_custom_typeshed(typeshed: &str) -> TestDb {
+        let db = TestDb::new();
+
+        Program::new(
+            &db,
+            TargetVersion::Py38,
+            SearchPathSettings {
+                extra_paths: Vec::new(),
+                workspace_root: SystemPathBuf::from("/src"),
+                site_packages: None,
+                custom_typeshed: Some(SystemPathBuf::from(typeshed)),
+            },
+        );
+
+        db
+    }
+
     fn assert_public_ty(db: &TestDb, file_name: &str, symbol_name: &str, expected: &str) {
         let file = system_path_to_file(db, file_name).expect("Expected file to exist.");
 
@@ -1370,6 +1400,80 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn builtin_symbol_vendored_stdlib() -> anyhow::Result<()> {
+        let mut db = setup_db();
+
+        db.write_file("/src/a.py", "c = copyright")?;
+
+        assert_public_ty(&db, "/src/a.py", "c", "Literal[copyright]");
+
+        Ok(())
+    }
+
+    #[test]
+    fn builtin_symbol_custom_stdlib() -> anyhow::Result<()> {
+        let mut db = setup_db_with_custom_typeshed("/typeshed");
+
+        db.write_files([
+            ("/src/a.py", "c = copyright"),
+            (
+                "/typeshed/stdlib/builtins.pyi",
+                "def copyright() -> None: ...",
+            ),
+            ("/typeshed/stdlib/VERSIONS", "builtins: 3.8-"),
+        ])?;
+
+        assert_public_ty(&db, "/src/a.py", "c", "Literal[copyright]");
+
+        Ok(())
+    }
+
+    #[test]
+    fn unknown_global_later_defined() -> anyhow::Result<()> {
+        let mut db = setup_db();
+
+        db.write_file("/src/a.py", "x = foo; foo = 1")?;
+
+        assert_public_ty(&db, "/src/a.py", "x", "Unbound");
+
+        Ok(())
+    }
+
+    #[test]
+    fn unknown_builtin_later_defined() -> anyhow::Result<()> {
+        let mut db = setup_db_with_custom_typeshed("/typeshed");
+
+        db.write_files([
+            ("/src/a.py", "x = foo"),
+            ("/typeshed/stdlib/builtins.pyi", "foo = bar; bar = 1"),
+            ("/typeshed/stdlib/VERSIONS", "builtins: 3.8-"),
+        ])?;
+
+        assert_public_ty(&db, "/src/a.py", "x", "Unbound");
+
+        Ok(())
+    }
+
+    #[test]
+    fn import_builtins() -> anyhow::Result<()> {
+        let mut db = setup_db();
+
+        db.write_file("/src/a.py", "import builtins; x = builtins.copyright")?;
+
+        assert_public_ty(&db, "/src/a.py", "x", "Literal[copyright]");
+        // imported builtins module is the same file as the implicit builtins
+        let file = system_path_to_file(&db, "/src/a.py").expect("Expected file to exist.");
+        let builtins_ty = global_symbol_ty_by_name(&db, file, "builtins");
+        let Type::Module(builtins_file) = builtins_ty else {
+            panic!("Builtins are not a module?");
+        };
+        let implicit_builtins_file = builtins_scope(&db).expect("builtins to exist").file(&db);
+        assert_eq!(builtins_file, implicit_builtins_file);
+
+        Ok(())
+    }
+
     fn first_public_def<'db>(db: &'db TestDb, file: File, name: &str) -> Definition<'db> {
         let scope = global_scope(db, file);
         *use_def_map(db, scope)