[red-knot] Add type inference for basic for loops (#13195)

Author: AlexWaygood

crates/red_knot_python_semantic/src/builtins.rs

@@ -10,7 +10,6 @@ pub use python_version::PythonVersion;
 pub use semantic_model::{HasTy, SemanticModel};
 
 pub mod ast_node_ref;
-mod builtins;
 mod db;
 mod module_name;
 mod module_resolver;
@@ -20,6 +19,7 @@ mod python_version;
 pub mod semantic_index;
 mod semantic_model;
 pub(crate) mod site_packages;
+mod stdlib;
 pub mod types;
 
 type FxOrderSet<V> = ordermap::set::OrderSet<V, BuildHasherDefault<FxHasher>>;

crates/red_knot_python_semantic/src/lib.rs

@@ -8,7 +8,7 @@ use crate::module_name::ModuleName;
 use crate::module_resolver::{resolve_module, Module};
 use crate::semantic_index::ast_ids::HasScopedAstId;
 use crate::semantic_index::semantic_index;
-use crate::types::{definition_ty, global_symbol_ty_by_name, infer_scope_types, Type};
+use crate::types::{definition_ty, global_symbol_ty, infer_scope_types, Type};
 use crate::Db;
 
 pub struct SemanticModel<'db> {
@@ -40,7 +40,7 @@ impl<'db> SemanticModel<'db> {
     }
 
     pub fn global_symbol_ty(&self, module: &Module, symbol_name: &str) -> Type<'db> {
-        global_symbol_ty_by_name(self.db, module.file(), symbol_name)
+        global_symbol_ty(self.db, module.file(), symbol_name)
     }
 }
 

crates/red_knot_python_semantic/src/semantic_model.rs

@@ -0,0 +1,77 @@
+use crate::module_name::ModuleName;
+use crate::module_resolver::resolve_module;
+use crate::semantic_index::global_scope;
+use crate::semantic_index::symbol::ScopeId;
+use crate::types::{global_symbol_ty, Type};
+use crate::Db;
+
+/// Enumeration of various core stdlib modules, for which we have dedicated Salsa queries.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+enum CoreStdlibModule {
+    Builtins,
+    Types,
+    Typeshed,
+}
+
+impl CoreStdlibModule {
+    fn name(self) -> ModuleName {
+        let module_name = match self {
+            Self::Builtins => "builtins",
+            Self::Types => "types",
+            Self::Typeshed => "_typeshed",
+        };
+        ModuleName::new_static(module_name)
+            .unwrap_or_else(|| panic!("{module_name} should be a valid module name!"))
+    }
+}
+
+/// Lookup the type of `symbol` in a given core module
+///
+/// Returns `Unbound` if the given core module cannot be resolved for some reason
+fn core_module_symbol_ty<'db>(
+    db: &'db dyn Db,
+    core_module: CoreStdlibModule,
+    symbol: &str,
+) -> Type<'db> {
+    resolve_module(db, core_module.name())
+        .map(|module| global_symbol_ty(db, module.file(), symbol))
+        .unwrap_or(Type::Unbound)
+}
+
+/// Lookup the type of `symbol` in the builtins namespace.
+///
+/// Returns `Unbound` if the `builtins` module isn't available for some reason.
+#[inline]
+pub(crate) fn builtins_symbol_ty<'db>(db: &'db dyn Db, symbol: &str) -> Type<'db> {
+    core_module_symbol_ty(db, CoreStdlibModule::Builtins, symbol)
+}
+
+/// Lookup the type of `symbol` in the `types` module namespace.
+///
+/// Returns `Unbound` if the `types` module isn't available for some reason.
+#[inline]
+pub(crate) fn types_symbol_ty<'db>(db: &'db dyn Db, symbol: &str) -> Type<'db> {
+    core_module_symbol_ty(db, CoreStdlibModule::Types, symbol)
+}
+
+/// Lookup the type of `symbol` in the `_typeshed` module namespace.
+///
+/// Returns `Unbound` if the `_typeshed` module isn't available for some reason.
+#[inline]
+pub(crate) fn typeshed_symbol_ty<'db>(db: &'db dyn Db, symbol: &str) -> Type<'db> {
+    core_module_symbol_ty(db, CoreStdlibModule::Typeshed, symbol)
+}
+
+/// Get the scope of a core stdlib module.
+///
+/// Can return `None` if a custom typeshed is used that is missing the core module in question.
+fn core_module_scope(db: &dyn Db, core_module: CoreStdlibModule) -> Option<ScopeId<'_>> {
+    resolve_module(db, core_module.name()).map(|module| global_scope(db, module.file()))
+}
+
+/// Get the `builtins` module scope.
+///
+/// Can return `None` if a custom typeshed is used that is missing `builtins.pyi`.
+pub(crate) fn builtins_module_scope(db: &dyn Db) -> Option<ScopeId<'_>> {
+    core_module_scope(db, CoreStdlibModule::Builtins)
+}