Bind top-most parent when importing nested module (astral-sh#14946)

When importing a nested module, we were correctly creating a binding for
the top-most parent, but we were binding that to the nested module, not
to that parent module. Moreover, we weren't treating those submodules as
members of their containing parents. This PR addresses both issues, so
that nested imports work as expected.

As discussed in ~Slack~ whatever chat app I find myself in these days
:smile:, this requires keeping track of which modules have been imported
within the current file, so that when we resolve member access on a
module reference, we can see if that member has been imported as a
submodule. If so, we return the submodule reference immediately, instead
of checking whether the parent module's definition defines the symbol.

This is currently done in a flow insensitive manner. The `SemanticIndex`
now tracks all of the modules that are imported (via `import`, not via
`from...import`). The member access logic mentioned above currently only
considers module imports in the file containing the attribute
expression.

---------

Co-authored-by: Carl Meyer <[email protected]>

crates/red_knot_python_semantic/resources/mdtest/import/basic.md

@@ -25,3 +25,69 @@ reveal_type(D)  # revealed: Literal[C]
 ```py path=b.py
 class C: ...
 ```
+
+## Nested
+
+```py
+import a.b
+
+reveal_type(a.b.C)  # revealed: Literal[C]
+```
+
+```py path=a/__init__.py
+```
+
+```py path=a/b.py
+class C: ...
+```
+
+## Deeply nested
+
+```py
+import a.b.c
+
+reveal_type(a.b.c.C)  # revealed: Literal[C]
+```
+
+```py path=a/__init__.py
+```
+
+```py path=a/b/__init__.py
+```
+
+```py path=a/b/c.py
+class C: ...
+```
+
+## Nested with rename
+
+```py
+import a.b as b
+
+reveal_type(b.C)  # revealed: Literal[C]
+```
+
+```py path=a/__init__.py
+```
+
+```py path=a/b.py
+class C: ...
+```
+
+## Deeply nested with rename
+
+```py
+import a.b.c as c
+
+reveal_type(c.C)  # revealed: Literal[C]
+```
+
+```py path=a/__init__.py
+```
+
+```py path=a/b/__init__.py
+```
+
+```py path=a/b/c.py
+class C: ...
+```

crates/red_knot_python_semantic/resources/mdtest/import/tracking.md

@@ -0,0 +1,100 @@
+# Tracking imported modules
+
+These tests depend on how we track which modules have been imported. There are currently two
+characteristics of our module tracking that can lead to inaccuracies:
+
+- Imports are tracked on a per-file basis. At runtime, importing a submodule in one file makes that
+    submodule globally available via any reference to the containing package. We will flag an error
+    if a file tries to access a submodule without there being an import of that submodule _in that
+    same file_.
+
+    This is a purposeful decision, and not one we plan to change. If a module wants to re-export some
+    other module that it imports, there are ways to do that (tested below) that are blessed by the
+    typing spec and that are visible to our file-scoped import tracking.
+
+- Imports are tracked flow-insensitively: submodule accesses are allowed and resolved if that
+    submodule is imported _anywhere in the file_. This handles the common case where all imports are
+    grouped at the top of the file, and is easiest to implement. We might revisit this decision and
+    track submodule imports flow-sensitively, in which case we will have to update the assertions in
+    some of these tests.
+
+## Import submodule later in file
+
+This test highlights our flow-insensitive analysis, since we access the `a.b` submodule before it
+has been imported.
+
+```py
+import a
+
+# Would be an error with flow-sensitive tracking
+reveal_type(a.b.C)  # revealed: Literal[C]
+
+import a.b
+```
+
+```py path=a/__init__.py
+```
+
+```py path=a/b.py
+class C: ...
+```
+
+## Rename a re-export
+
+This test highlights how import tracking is local to each file, but specifically to the file where a
+containing module is first referenced. This allows the main module to see that `q.a` contains a
+submodule `b`, even though `a.b` is never imported in the main module.
+
+```py
+from q import a, b
+
+reveal_type(b)  # revealed: <module 'a.b'>
+reveal_type(b.C)  # revealed: Literal[C]
+
+reveal_type(a.b)  # revealed: <module 'a.b'>
+reveal_type(a.b.C)  # revealed: Literal[C]
+```
+
+```py path=a/__init__.py
+```
+
+```py path=a/b.py
+class C: ...
+```
+
+```py path=q.py
+import a as a
+import a.b as b
+```
+
+## Attribute overrides submodule
+
+Technically, either a submodule or a non-module attribute could shadow the other, depending on the
+ordering of when the submodule is loaded relative to the parent module's `__init__.py` file being
+evaluated. We have chosen to always have the submodule take priority. (This matches pyright's
+current behavior, and opposite of mypy's current behavior.)
+
+```py
+import sub.b
+import attr.b
+
+# In the Python interpreter, `attr.b` is Literal[1]
+reveal_type(sub.b)  # revealed: <module 'sub.b'>
+reveal_type(attr.b)  # revealed: <module 'attr.b'>
+```
+
+```py path=sub/__init__.py
+b = 1
+```
+
+```py path=sub/b.py
+```
+
+```py path=attr/__init__.py
+from . import b as _
+
+b = 1
+```
+
+```py path=attr/b.py
+```

crates/red_knot_python_semantic/resources/mdtest/type_of/basic.md

@@ -61,10 +61,8 @@ class B: ...
 ```py path=a/test.py
 import a.b
 
-# TODO: no diagnostic
-# error: [unresolved-attribute]
 def f(c: type[a.b.C]):
-    reveal_type(c)  # revealed: @Todo(unsupported type[X] special form)
+    reveal_type(c)  # revealed: type[C]
 ```
 
 ```py path=a/__init__.py

crates/red_knot_python_semantic/src/module_name.rs

@@ -186,6 +186,26 @@ impl ModuleName {
         self.0.push('.');
         self.0.push_str(other);
     }
+
+    /// Returns an iterator of this module name and all of its parent modules.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use red_knot_python_semantic::ModuleName;
+    ///
+    /// assert_eq!(
+    ///     ModuleName::new_static("foo.bar.baz").unwrap().ancestors().collect::<Vec<_>>(),
+    ///     vec![
+    ///         ModuleName::new_static("foo.bar.baz").unwrap(),
+    ///         ModuleName::new_static("foo.bar").unwrap(),
+    ///         ModuleName::new_static("foo").unwrap(),
+    ///     ],
+    /// );
+    /// ```
+    pub fn ancestors(&self) -> impl Iterator<Item = Self> {
+        std::iter::successors(Some(self.clone()), Self::parent)
+    }
 }
 
 impl Deref for ModuleName {

crates/red_knot_python_semantic/src/module_resolver/module.rs

@@ -7,7 +7,7 @@ use super::path::SearchPath;
 use crate::module_name::ModuleName;
 
 /// Representation of a Python module.
-#[derive(Clone, PartialEq, Eq)]
+#[derive(Clone, PartialEq, Eq, Hash)]
 pub struct Module {
     inner: Arc<ModuleInner>,
 }
@@ -61,7 +61,7 @@ impl std::fmt::Debug for Module {
     }
 }
 
-#[derive(PartialEq, Eq)]
+#[derive(PartialEq, Eq, Hash)]
 struct ModuleInner {
     name: ModuleName,
     kind: ModuleKind,

crates/red_knot_python_semantic/src/semantic_index.rs

@@ -1,13 +1,14 @@
 use std::iter::FusedIterator;
 use std::sync::Arc;
 
-use rustc_hash::{FxBuildHasher, FxHashMap};
+use rustc_hash::{FxBuildHasher, FxHashMap, FxHashSet};
 use salsa::plumbing::AsId;
 
 use ruff_db::files::File;
 use ruff_db::parsed::parsed_module;
 use ruff_index::{IndexSlice, IndexVec};
 
+use crate::module_name::ModuleName;
 use crate::semantic_index::ast_ids::node_key::ExpressionNodeKey;
 use crate::semantic_index::ast_ids::AstIds;
 use crate::semantic_index::builder::SemanticIndexBuilder;
@@ -60,6 +61,12 @@ pub(crate) fn symbol_table<'db>(db: &'db dyn Db, scope: ScopeId<'db>) -> Arc<Sym
     index.symbol_table(scope.file_scope_id(db))
 }
 
+/// Returns the set of modules that are imported anywhere in `file`.
+#[salsa::tracked]
+pub(crate) fn imported_modules<'db>(db: &'db dyn Db, file: File) -> Arc<FxHashSet<ModuleName>> {
+    semantic_index(db, file).imported_modules.clone()
+}
+
 /// Returns the use-def map for a specific `scope`.
 ///
 /// Using [`use_def_map`] over [`semantic_index`] has the advantage that
@@ -116,6 +123,9 @@ pub(crate) struct SemanticIndex<'db> {
     /// changing a file invalidates all dependents.
     ast_ids: IndexVec<FileScopeId, AstIds>,
 
+    /// The set of modules that are imported anywhere within this file.
+    imported_modules: Arc<FxHashSet<ModuleName>>,
+
     /// Flags about the global scope (code usage impacting inference)
     has_future_annotations: bool,
 }

crates/red_knot_python_semantic/src/semantic_index/builder.rs

@@ -1,7 +1,7 @@
 use std::sync::Arc;
 
 use except_handlers::TryNodeContextStackManager;
-use rustc_hash::FxHashMap;
+use rustc_hash::{FxHashMap, FxHashSet};
 
 use ruff_db::files::File;
 use ruff_db::parsed::ParsedModule;
@@ -12,6 +12,7 @@ use ruff_python_ast::visitor::{walk_expr, walk_pattern, walk_stmt, Visitor};
 use ruff_python_ast::{BoolOp, Expr};
 
 use crate::ast_node_ref::AstNodeRef;
+use crate::module_name::ModuleName;
 use crate::semantic_index::ast_ids::node_key::ExpressionNodeKey;
 use crate::semantic_index::ast_ids::AstIdsBuilder;
 use crate::semantic_index::definition::{
@@ -79,6 +80,7 @@ pub(super) struct SemanticIndexBuilder<'db> {
     scopes_by_expression: FxHashMap<ExpressionNodeKey, FileScopeId>,
     definitions_by_node: FxHashMap<DefinitionNodeKey, Definition<'db>>,
     expressions_by_node: FxHashMap<ExpressionNodeKey, Expression<'db>>,
+    imported_modules: FxHashSet<ModuleName>,
 }
 
 impl<'db> SemanticIndexBuilder<'db> {
@@ -105,6 +107,8 @@ impl<'db> SemanticIndexBuilder<'db> {
             scopes_by_node: FxHashMap::default(),
             definitions_by_node: FxHashMap::default(),
             expressions_by_node: FxHashMap::default(),
+
+            imported_modules: FxHashSet::default(),
         };
 
         builder.push_scope_with_parent(NodeWithScopeRef::Module, None);
@@ -558,6 +562,7 @@ impl<'db> SemanticIndexBuilder<'db> {
             scopes_by_expression: self.scopes_by_expression,
             scopes_by_node: self.scopes_by_node,
             use_def_maps,
+            imported_modules: Arc::new(self.imported_modules),
             has_future_annotations: self.has_future_annotations,
         }
     }
@@ -661,6 +666,12 @@ where
             }
             ast::Stmt::Import(node) => {
                 for alias in &node.names {
+                    // Mark the imported module, and all of its parents, as being imported in this
+                    // file.
+                    if let Some(module_name) = ModuleName::new(&alias.name) {
+                        self.imported_modules.extend(module_name.ancestors());
+                    }
+
                     let symbol_name = if let Some(asname) = &alias.asname {
                         asname.id.clone()
                     } else {