[ty] Support custom builtins (#22021)

Author: Wizzerinus

crates/ty_python_semantic/resources/mdtest/import/builtins.md

@@ -76,3 +76,63 @@ def reveal_type(obj, /): ...
 ```py
 reveal_type(foo)  # revealed: Unknown
 ```
+
+## Builtins imported from custom project-level stubs
+
+The project can add or replace builtins with the `__builtins__.pyi` stub. They will take precedence
+over the typeshed ones.
+
+```py
+reveal_type(foo)  # revealed: int
+reveal_type(bar)  # revealed: str
+reveal_type(quux(1))  # revealed: int
+b = baz  # error: [unresolved-reference]
+
+reveal_type(ord(100))  # revealed: bool
+a = ord("a")  # error: [invalid-argument-type]
+
+bar = int(123)
+reveal_type(bar)  # revealed: int
+```
+
+`__builtins__.pyi`:
+
+```pyi
+foo: int = ...
+bar: str = ...
+
+def quux(value: int) -> int: ...
+
+unused: str = ...
+
+def ord(x: int) -> bool: ...
+```
+
+Builtins stubs are searched relative to the project root, not the file using them.
+
+`under/some/folder.py`:
+
+```py
+reveal_type(foo)  # revealed: int
+reveal_type(bar)  # revealed: str
+```
+
+## Assigning custom builtins
+
+```py
+import builtins
+
+builtins.foo = 123
+builtins.bar = 456  # error: [unresolved-attribute]
+builtins.baz = 789  # error: [invalid-assignment]
+builtins.chr = lambda x: str(x)  # error: [invalid-assignment]
+builtins.chr = 10
+```
+
+`__builtins__.pyi`:
+
+```pyi
+foo: int
+baz: str
+chr: int
+```

crates/ty_python_semantic/src/place.rs

@@ -1,6 +1,8 @@
 use ruff_db::files::File;
 use ruff_python_ast::PythonVersion;
-use ty_module_resolver::{KnownModule, file_to_module, resolve_module_confident};
+use ty_module_resolver::{
+    KnownModule, Module, ModuleName, file_to_module, resolve_module_confident,
+};
 
 use crate::dunder_all::dunder_all_names;
 use crate::semantic_index::definition::{Definition, DefinitionState};
@@ -380,25 +382,29 @@ pub(crate) fn imported_symbol<'db>(
 /// and should not be used when a symbol is being explicitly imported from the `builtins` module
 /// (e.g. `from builtins import int`).
 pub(crate) fn builtins_symbol<'db>(db: &'db dyn Db, symbol: &str) -> PlaceAndQualifiers<'db> {
-    resolve_module_confident(db, &KnownModule::Builtins.name())
-        .and_then(|module| {
-            let file = module.file(db)?;
-            Some(
-                symbol_impl(
-                    db,
-                    global_scope(db, file),
-                    symbol,
-                    RequiresExplicitReExport::Yes,
-                    ConsideredDefinitions::EndOfScope,
-                )
-                .or_fall_back_to(db, || {
-                    // We're looking up in the builtins namespace and not the module, so we should
-                    // do the normal lookup in `types.ModuleType` and not the special one as in
-                    // `imported_symbol`.
-                    module_type_implicit_global_symbol(db, symbol)
-                }),
-            )
-        })
+    let resolver = |module: Module<'_>| {
+        let file = module.file(db)?;
+        let found_symbol = symbol_impl(
+            db,
+            global_scope(db, file),
+            symbol,
+            RequiresExplicitReExport::Yes,
+            ConsideredDefinitions::EndOfScope,
+        )
+        .or_fall_back_to(db, || {
+            // We're looking up in the builtins namespace and not the module, so we should
+            // do the normal lookup in `types.ModuleType` and not the special one as in
+            // `imported_symbol`.
+            module_type_implicit_global_symbol(db, symbol)
+        });
+        // If this symbol is not present in project-level builtins, search in the default ones.
+        found_symbol
+            .ignore_possibly_undefined()
+            .map(|_| found_symbol)
+    };
+    resolve_module_confident(db, &ModuleName::new_static("__builtins__").unwrap())
+        .and_then(&resolver)
+        .or_else(|| resolve_module_confident(db, &KnownModule::Builtins.name()).and_then(resolver))
         .unwrap_or_default()
 }
 

crates/ty_python_semantic/src/types/infer/builder.rs

@@ -4963,7 +4963,16 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
             }
 
             Type::ModuleLiteral(module) => {
-                if let Place::Defined(attr_ty, _, _) = module.static_member(db, attribute).place {
+                let sym = if module
+                    .module(db)
+                    .known(db)
+                    .is_some_and(KnownModule::is_builtins)
+                {
+                    builtins_symbol(db, attribute)
+                } else {
+                    module.static_member(db, attribute)
+                };
+                if let Place::Defined(attr_ty, _, _) = sym.place {
                     let value_ty = infer_value_ty(self, TypeContext::new(Some(attr_ty)));
 
                     let assignable = value_ty.is_assignable_to(db, attr_ty);