Refactor symbol lookup APIs to hide re-export implementation details (#16133)

## Summary

This PR refactors the symbol lookup APIs to better facilitate the
re-export implementation. Specifically,
* Add `module_type_symbol` which returns the `Symbol` that's a member of
`types.ModuleType`
* Rename `symbol` -> `symbol_impl`; add `symbol` which delegates to
`symbol_impl` with `RequireExplicitReExport::No`
* Update `global_symbol` to do `symbol_impl` -> fall back to
`module_type_symbol` and default to `RequireExplicitReExport::No`
* Add `imported_symbol` to do `symbol_impl` with
`RequireExplicitReExport` as `Yes` if the module is in a stub file else
`No`
* Update `known_module_symbol` to use `imported_symbol` with a fallback
to `module_type_symbol`
* Update `ModuleLiteralType::member` to use `imported_symbol` with a
custom fallback

We could potentially also update `symbol_from_declarations` and
`symbol_from_bindings` to avoid passing in the `RequireExplicitReExport`
as it would be always `No` if called directly. We could add
`symbol_from_declarations_impl` and `symbol_from_bindings_impl`.

Looking at the `_impl` functions, I think we should move all of these
symbol related logic into `symbol.rs` where `Symbol` is defined and the
`_impl` could be private while we expose the public APIs at the crate
level. This would also make the `RequireExplicitReExport` an
implementation detail and the caller doesn't need to worry about it.

Author: dhruvmanila

crates/red_knot_python_semantic/src/semantic_index/definition.rs

@@ -50,10 +50,6 @@ impl<'db> Definition<'db> {
         self.kind(db).category()
     }
 
-    pub(crate) fn in_stub(self, db: &'db dyn Db) -> bool {
-        self.file(db).is_stub(db.upcast())
-    }
-
     pub(crate) fn is_declaration(self, db: &'db dyn Db) -> bool {
         self.kind(db).category().is_declaration()
     }

crates/red_knot_python_semantic/src/stdlib.rs

@@ -2,7 +2,7 @@ use crate::module_resolver::{resolve_module, KnownModule};
 use crate::semantic_index::global_scope;
 use crate::semantic_index::symbol::ScopeId;
 use crate::symbol::Symbol;
-use crate::types::{global_symbol, SymbolLookup};
+use crate::types::imported_symbol;
 use crate::Db;
 
 /// Lookup the type of `symbol` in a given known module
@@ -14,18 +14,10 @@ pub(crate) fn known_module_symbol<'db>(
     symbol: &str,
 ) -> Symbol<'db> {
     resolve_module(db, &known_module.name())
-        .map(|module| global_symbol(db, SymbolLookup::External, module.file(), symbol))
+        .map(|module| imported_symbol(db, &module, symbol))
         .unwrap_or(Symbol::Unbound)
 }
 
-/// Lookup the type of `symbol` in the builtins namespace.
-///
-/// Returns `Symbol::Unbound` if the `builtins` module isn't available for some reason.
-#[inline]
-pub(crate) fn builtins_symbol<'db>(db: &'db dyn Db, symbol: &str) -> Symbol<'db> {
-    known_module_symbol(db, KnownModule::Builtins, symbol)
-}
-
 /// Lookup the type of `symbol` in the `typing` module namespace.
 ///
 /// Returns `Symbol::Unbound` if the `typing` module isn't available for some reason.

crates/red_knot_python_semantic/src/types.rs

@@ -32,7 +32,7 @@ use crate::semantic_index::{
     use_def_map, BindingWithConstraints, BindingWithConstraintsIterator, DeclarationWithConstraint,
     DeclarationsIterator,
 };
-use crate::stdlib::{builtins_symbol, known_module_symbol, typing_extensions_symbol};
+use crate::stdlib::{known_module_symbol, typing_extensions_symbol};
 use crate::suppression::check_suppressions;
 use crate::symbol::{Boundness, Symbol};
 use crate::types::call::{
@@ -107,40 +107,37 @@ fn widen_type_for_undeclared_public_symbol<'db>(
 }
 
 #[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
-pub(crate) enum SymbolLookup {
-    /// Look up the symbol as seen from within the same module.
-    Internal,
-    /// Look up the symbol as seen from outside the module.
-    External,
+enum RequiresExplicitReExport {
+    Yes,
+    No,
 }
 
-impl SymbolLookup {
-    const fn is_external(self) -> bool {
-        matches!(self, Self::External)
+impl RequiresExplicitReExport {
+    const fn is_yes(self) -> bool {
+        matches!(self, RequiresExplicitReExport::Yes)
     }
 }
 
-/// Infer the public type of a symbol (its type as seen from outside its scope).
-fn symbol<'db>(
+fn symbol_impl<'db>(
     db: &'db dyn Db,
-    lookup: SymbolLookup,
     scope: ScopeId<'db>,
     name: &str,
+    requires_explicit_reexport: RequiresExplicitReExport,
 ) -> Symbol<'db> {
     #[salsa::tracked]
     fn symbol_by_id<'db>(
         db: &'db dyn Db,
-        lookup: SymbolLookup,
         scope: ScopeId<'db>,
         symbol_id: ScopedSymbolId,
+        requires_explicit_reexport: RequiresExplicitReExport,
     ) -> Symbol<'db> {
         let use_def = use_def_map(db, scope);
 
         // If the symbol is declared, the public type is based on declarations; otherwise, it's based
         // on inference from bindings.
 
         let declarations = use_def.public_declarations(symbol_id);
-        let declared = symbol_from_declarations(db, lookup, declarations);
+        let declared = symbol_from_declarations(db, declarations, requires_explicit_reexport);
         let is_final = declared.as_ref().is_ok_and(SymbolAndQualifiers::is_final);
         let declared = declared.map(|SymbolAndQualifiers(symbol, _)| symbol);
 
@@ -150,7 +147,7 @@ fn symbol<'db>(
             // Symbol is possibly declared
             Ok(Symbol::Type(declared_ty, Boundness::PossiblyUnbound)) => {
                 let bindings = use_def.public_bindings(symbol_id);
-                let inferred = symbol_from_bindings(db, lookup, bindings);
+                let inferred = symbol_from_bindings(db, bindings, requires_explicit_reexport);
 
                 match inferred {
                     // Symbol is possibly undeclared and definitely unbound
@@ -170,7 +167,7 @@ fn symbol<'db>(
             // Symbol is undeclared, return the union of `Unknown` with the inferred type
             Ok(Symbol::Unbound) => {
                 let bindings = use_def.public_bindings(symbol_id);
-                let inferred = symbol_from_bindings(db, lookup, bindings);
+                let inferred = symbol_from_bindings(db, bindings, requires_explicit_reexport);
 
                 // `__slots__` is a symbol with special behavior in Python's runtime. It can be
                 // modified externally, but those changes do not take effect. We therefore issue
@@ -232,7 +229,7 @@ fn symbol<'db>(
 
     symbol_table(db, scope)
         .symbol_id_by_name(name)
-        .map(|symbol| symbol_by_id(db, lookup, scope, symbol))
+        .map(|symbol| symbol_by_id(db, scope, symbol, requires_explicit_reexport))
         .unwrap_or(Symbol::Unbound)
 }
 
@@ -271,27 +268,99 @@ fn module_type_symbols<'db>(db: &'db dyn Db) -> smallvec::SmallVec<[ast::name::N
         .collect()
 }
 
-pub(crate) fn global_symbol<'db>(
-    db: &'db dyn Db,
-    lookup: SymbolLookup,
-    file: File,
-    name: &str,
-) -> Symbol<'db> {
-    // Not defined explicitly in the global scope?
-    // All modules are instances of `types.ModuleType`;
-    // look it up there (with a few very special exceptions)
-    symbol(db, lookup, global_scope(db, file), name).or_fall_back_to(db, || {
-        if module_type_symbols(db)
-            .iter()
-            .any(|module_type_member| &**module_type_member == name)
-        {
-            KnownClass::ModuleType.to_instance(db).member(db, name)
-        } else {
+/// Return the symbol for a member of `types.ModuleType`.
+pub(crate) fn module_type_symbol<'db>(db: &'db dyn Db, name: &str) -> Symbol<'db> {
+    if module_type_symbols(db)
+        .iter()
+        .any(|module_type_member| &**module_type_member == name)
+    {
+        KnownClass::ModuleType.to_instance(db).member(db, name)
+    } else {
+        Symbol::Unbound
+    }
+}
+
+/// Infer the public type of a symbol (its type as seen from outside its scope) in the given
+/// `scope`.
+fn symbol<'db>(db: &'db dyn Db, scope: ScopeId<'db>, name: &str) -> Symbol<'db> {
+    symbol_impl(db, scope, name, RequiresExplicitReExport::No)
+}
+
+/// Infers the public type of a module-global symbol as seen from within the same file.
+///
+/// If it's not defined explicitly in the global scope, it will look it up in `types.ModuleType`
+/// with a few very special exceptions.
+///
+/// Use [`imported_symbol`] to perform the lookup as seen from outside the file (e.g. via imports).
+pub(crate) fn global_symbol<'db>(db: &'db dyn Db, file: File, name: &str) -> Symbol<'db> {
+    symbol_impl(
+        db,
+        global_scope(db, file),
+        name,
+        RequiresExplicitReExport::No,
+    )
+    .or_fall_back_to(db, || module_type_symbol(db, name))
+}
+
+/// Infers the public type of an imported symbol.
+pub(crate) fn imported_symbol<'db>(db: &'db dyn Db, module: &Module, name: &str) -> Symbol<'db> {
+    // If it's not found in the global scope, check if it's present as an instance on
+    // `types.ModuleType` or `builtins.object`.
+    //
+    // We do a more limited version of this in `global_symbol`, but there are two crucial
+    // differences here:
+    // - If a member is looked up as an attribute, `__init__` is also available on the module, but
+    //   it isn't available as a global from inside the module
+    // - If a member is looked up as an attribute, members on `builtins.object` are also available
+    //   (because `types.ModuleType` inherits from `object`); these attributes are also not
+    //   available as globals from inside the module.
+    //
+    // The same way as in `global_symbol`, however, we need to be careful to ignore
+    // `__getattr__`. Typeshed has a fake `__getattr__` on `types.ModuleType` to help out with
+    // dynamic imports; we shouldn't use it for `ModuleLiteral` types where we know exactly which
+    // module we're dealing with.
+    external_symbol_impl(db, module.file(), name).or_fall_back_to(db, || {
+        if name == "__getattr__" {
             Symbol::Unbound
+        } else {
+            KnownClass::ModuleType.to_instance(db).member(db, name)
         }
     })
 }
 
+/// Lookup the type of `symbol` in the builtins namespace.
+///
+/// Returns `Symbol::Unbound` if the `builtins` module isn't available for some reason.
+///
+/// Note that this function is only intended for use in the context of the builtins *namespace*
+/// 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) -> Symbol<'db> {
+    resolve_module(db, &KnownModule::Builtins.name())
+        .map(|module| {
+            external_symbol_impl(db, module.file(), symbol).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_symbol(db, symbol)
+            })
+        })
+        .unwrap_or(Symbol::Unbound)
+}
+
+fn external_symbol_impl<'db>(db: &'db dyn Db, file: File, name: &str) -> Symbol<'db> {
+    symbol_impl(
+        db,
+        global_scope(db, file),
+        name,
+        if file.is_stub(db.upcast()) {
+            RequiresExplicitReExport::Yes
+        } else {
+            RequiresExplicitReExport::No
+        },
+    )
+}
+
 /// Infer the type of a binding.
 pub(crate) fn binding_type<'db>(db: &'db dyn Db, definition: Definition<'db>) -> Type<'db> {
     let inference = infer_definition_types(db, definition);
@@ -340,14 +409,14 @@ fn definition_expression_type<'db>(
 /// The type will be a union if there are multiple bindings with different types.
 fn symbol_from_bindings<'db>(
     db: &'db dyn Db,
-    lookup: SymbolLookup,
     bindings_with_constraints: BindingWithConstraintsIterator<'_, 'db>,
+    requires_explicit_reexport: RequiresExplicitReExport,
 ) -> Symbol<'db> {
     let visibility_constraints = bindings_with_constraints.visibility_constraints;
     let mut bindings_with_constraints = bindings_with_constraints.peekable();
 
     let is_non_exported = |binding: Definition<'db>| {
-        lookup.is_external() && !binding.is_reexported(db) && binding.in_stub(db)
+        requires_explicit_reexport.is_yes() && !binding.is_reexported(db)
     };
 
     let unbound_visibility = match bindings_with_constraints.peek() {
@@ -471,14 +540,14 @@ type SymbolFromDeclarationsResult<'db> =
 /// [`TypeQualifiers`] that have been specified on the declaration(s).
 fn symbol_from_declarations<'db>(
     db: &'db dyn Db,
-    lookup: SymbolLookup,
     declarations: DeclarationsIterator<'_, 'db>,
+    requires_explicit_reexport: RequiresExplicitReExport,
 ) -> SymbolFromDeclarationsResult<'db> {
     let visibility_constraints = declarations.visibility_constraints;
     let mut declarations = declarations.peekable();
 
     let is_non_exported = |declaration: Definition<'db>| {
-        lookup.is_external() && !declaration.is_reexported(db) && declaration.in_stub(db)
+        requires_explicit_reexport.is_yes() && !declaration.is_reexported(db)
     };
 
     let undeclared_visibility = match declarations.peek() {
@@ -3839,31 +3908,7 @@ impl<'db> ModuleLiteralType<'db> {
             }
         }
 
-        // If it's not found in the global scope, check if it's present as an instance
-        // on `types.ModuleType` or `builtins.object`.
-        //
-        // We do a more limited version of this in `global_symbol_ty`,
-        // but there are two crucial differences here:
-        // - If a member is looked up as an attribute, `__init__` is also available
-        //   on the module, but it isn't available as a global from inside the module
-        // - If a member is looked up as an attribute, members on `builtins.object`
-        //   are also available (because `types.ModuleType` inherits from `object`);
-        //   these attributes are also not available as globals from inside the module.
-        //
-        // The same way as in `global_symbol_ty`, however, we need to be careful to
-        // ignore `__getattr__`. Typeshed has a fake `__getattr__` on `types.ModuleType`
-        // to help out with dynamic imports; we shouldn't use it for `ModuleLiteral` types
-        // where we know exactly which module we're dealing with.
-        global_symbol(db, SymbolLookup::External, self.module(db).file(), name).or_fall_back_to(
-            db,
-            || {
-                if name == "__getattr__" {
-                    Symbol::Unbound
-                } else {
-                    KnownClass::ModuleType.to_instance(db).member(db, name)
-                }
-            },
-        )
+        imported_symbol(db, &self.module(db), name)
     }
 }
 
@@ -4198,7 +4243,7 @@ impl<'db> Class<'db> {
     /// traverse through the MRO until it finds the member.
     pub(crate) fn own_class_member(self, db: &'db dyn Db, name: &str) -> Symbol<'db> {
         let scope = self.body_scope(db);
-        symbol(db, SymbolLookup::Internal, scope, name)
+        symbol(db, scope, name)
     }
 
     /// Returns the `name` attribute of an instance of this class.
@@ -4340,7 +4385,7 @@ impl<'db> Class<'db> {
 
             let declarations = use_def.public_declarations(symbol_id);
 
-            match symbol_from_declarations(db, SymbolLookup::Internal, declarations) {
+            match symbol_from_declarations(db, declarations, RequiresExplicitReExport::No) {
                 Ok(SymbolAndQualifiers(Symbol::Type(declared_ty, _), qualifiers)) => {
                     // The attribute is declared in the class body.
 
@@ -4362,7 +4407,7 @@ impl<'db> Class<'db> {
                     // in a method, and it could also be *bound* in the class body (and/or in a method).
 
                     let bindings = use_def.public_bindings(symbol_id);
-                    let inferred = symbol_from_bindings(db, SymbolLookup::Internal, bindings);
+                    let inferred = symbol_from_bindings(db, bindings, RequiresExplicitReExport::No);
                     let inferred_ty = inferred.ignore_possibly_unbound();
 
                     Self::implicit_instance_attribute(db, body_scope, name, inferred_ty).into()
@@ -4980,7 +5025,7 @@ pub(crate) mod tests {
         )?;
 
         let bar = system_path_to_file(&db, "src/bar.py")?;
-        let a = global_symbol(&db, SymbolLookup::Internal, bar, "a");
+        let a = global_symbol(&db, bar, "a");
 
         assert_eq!(
             a.expect_type(),
@@ -4999,7 +5044,7 @@ pub(crate) mod tests {
         )?;
         db.clear_salsa_events();
 
-        let a = global_symbol(&db, SymbolLookup::Internal, bar, "a");
+        let a = global_symbol(&db, bar, "a");
 
         assert_eq!(
             a.expect_type(),