Handle nested imports correctly in from ... import (#15026)

#14946 fixed our handling of nested imports with the `import` statement,
but didn't touch `from...import` statements.

cf
#14826 (comment)

Author: dcreager

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

@@ -7,3 +7,25 @@ from import bar  # error: [invalid-syntax]
 
 reveal_type(bar)  # revealed: Unknown
 ```
+
+## Invalid nested module import
+
+TODO: This is correctly flagged as an error, but we could clean up the diagnostics that we report.
+
+```py
+# TODO: No second diagnostic
+# error: [invalid-syntax] "Expected ',', found '.'"
+# error: [unresolved-import] "Module `a` has no member `c`"
+from a import b.c
+
+# TODO: Should these be inferred as Unknown?
+reveal_type(b)  # revealed: <module 'a.b'>
+reveal_type(b.c)  # revealed: Literal[1]
+```
+
+```py path=a/__init__.py
+```
+
+```py path=a/b.py
+c = 1
+```

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

@@ -121,23 +121,44 @@ X = 42
 ```
 
 ```py path=package/bar.py
-# TODO: support submodule imports
-from . import foo  # error: [unresolved-import]
-
-y = foo.X
+from . import foo
 
-# TODO: should be `Literal[42]`
-reveal_type(y)  # revealed: Unknown
+reveal_type(foo.X)  # revealed: Literal[42]
 ```
 
 ## Non-existent + bare to module
 
+This test verifies that we emit an error when we try to import a symbol that is neither a submodule
+nor an attribute of `package`.
+
 ```py path=package/__init__.py
 ```
 
 ```py path=package/bar.py
-# TODO: support submodule imports
 from . import foo  # error: [unresolved-import]
 
 reveal_type(foo)  # revealed: Unknown
 ```
+
+## Import submodule from self
+
+We don't currently consider `from...import` statements when building up the `imported_modules` set
+in the semantic index. When accessing an attribute of a module, we only consider it a potential
+submodule when that submodule name appears in the `imported_modules` set. That means that submodules
+that are imported via `from...import` are not visible to our type inference if you also access that
+submodule via the attribute on its parent package.
+
+```py path=package/__init__.py
+```
+
+```py path=package/foo.py
+X = 42
+```
+
+```py path=package/bar.py
+from . import foo
+import package
+
+# error: [unresolved-attribute] "Type `<module 'package'>` has no attribute `foo`"
+reveal_type(package.foo.X)  # revealed: Unknown
+```

crates/red_knot_python_semantic/src/semantic_index.rs

@@ -62,6 +62,16 @@ pub(crate) fn symbol_table<'db>(db: &'db dyn Db, scope: ScopeId<'db>) -> Arc<Sym
 }
 
 /// Returns the set of modules that are imported anywhere in `file`.
+///
+/// This set only considers `import` statements, not `from...import` statements, because:
+///
+///   - In `from foo import bar`, we cannot determine whether `foo.bar` is a submodule (and is
+///     therefore imported) without looking outside the content of this file.  (We could turn this
+///     into a _potentially_ imported modules set, but that would change how it's used in our type
+///     inference logic.)
+///
+///   - We cannot resolve relative imports (which aren't allowed in `import` statements) without
+///     knowing the name of the current module, and whether it's a package.
 #[salsa::tracked]
 pub(crate) fn imported_modules<'db>(db: &'db dyn Db, file: File) -> Arc<FxHashSet<ModuleName>> {
     semantic_index(db, file).imported_modules.clone()

crates/red_knot_python_semantic/src/types/infer.rs

@@ -865,6 +865,14 @@ impl<'db> TypeInferenceBuilder<'db> {
         self.types.bindings.insert(definition, inferred_ty);
     }
 
+    fn add_unknown_declaration_with_binding(
+        &mut self,
+        node: AnyNodeRef,
+        definition: Definition<'db>,
+    ) {
+        self.add_declaration_with_binding(node, definition, Type::Unknown, Type::Unknown);
+    }
+
     fn infer_module(&mut self, module: &ast::ModModule) {
         self.infer_body(&module.body);
     }
@@ -2121,24 +2129,14 @@ impl<'db> TypeInferenceBuilder<'db> {
         // The name of the module being imported
         let Some(full_module_name) = ModuleName::new(name) else {
             tracing::debug!("Failed to resolve import due to invalid syntax");
-            self.add_declaration_with_binding(
-                alias.into(),
-                definition,
-                Type::Unknown,
-                Type::Unknown,
-            );
+            self.add_unknown_declaration_with_binding(alias.into(), definition);
             return;
         };
 
         // Resolve the module being imported.
         let Some(full_module_ty) = self.module_ty_from_name(&full_module_name) else {
             self.diagnostics.add_unresolved_module(alias, 0, Some(name));
-            self.add_declaration_with_binding(
-                alias.into(),
-                definition,
-                Type::Unknown,
-                Type::Unknown,
-            );
+            self.add_unknown_declaration_with_binding(alias.into(), definition);
             return;
         };
 
@@ -2152,11 +2150,11 @@ impl<'db> TypeInferenceBuilder<'db> {
             // parent package of that module.
             let topmost_parent_name =
                 ModuleName::new(full_module_name.components().next().unwrap()).unwrap();
-            if let Some(topmost_parent_ty) = self.module_ty_from_name(&topmost_parent_name) {
-                topmost_parent_ty
-            } else {
-                Type::Unknown
-            }
+            let Some(topmost_parent_ty) = self.module_ty_from_name(&topmost_parent_name) else {
+                self.add_unknown_declaration_with_binding(alias.into(), definition);
+                return;
+            };
+            topmost_parent_ty
         } else {
             // If there's no `as` clause and the imported module isn't nested, then the imported
             // module _is_ what we bind into the current scope.
@@ -2243,12 +2241,6 @@ impl<'db> TypeInferenceBuilder<'db> {
         // TODO:
         // - Absolute `*` imports (`from collections import *`)
         // - Relative `*` imports (`from ...foo import *`)
-        // - Submodule imports (`from collections import abc`,
-        //   where `abc` is a submodule of the `collections` package)
-        //
-        // For the last item, see the currently skipped tests
-        // `follow_relative_import_bare_to_module()` and
-        // `follow_nonexistent_import_bare_to_module()`.
         let ast::StmtImportFrom { module, level, .. } = import_from;
         let module = module.as_deref();
 
@@ -2271,46 +2263,13 @@ impl<'db> TypeInferenceBuilder<'db> {
                 .ok_or(ModuleNameResolutionError::InvalidSyntax)
         };
 
-        let ty = match module_name {
-            Ok(module_name) => {
-                if let Some(module_ty) = self.module_ty_from_name(&module_name) {
-                    let ast::Alias {
-                        range: _,
-                        name,
-                        asname: _,
-                    } = alias;
-
-                    match module_ty.member(self.db, &ast::name::Name::new(&name.id)) {
-                        Symbol::Type(ty, boundness) => {
-                            if boundness == Boundness::PossiblyUnbound {
-                                self.diagnostics.add_lint(
-                                    &POSSIBLY_UNBOUND_IMPORT,
-                                    AnyNodeRef::Alias(alias),
-                                    format_args!("Member `{name}` of module `{module_name}` is possibly unbound", ),
-                                );
-                            }
-
-                            ty
-                        }
-                        Symbol::Unbound => {
-                            self.diagnostics.add_lint(
-                                &UNRESOLVED_IMPORT,
-                                AnyNodeRef::Alias(alias),
-                                format_args!("Module `{module_name}` has no member `{name}`",),
-                            );
-                            Type::Unknown
-                        }
-                    }
-                } else {
-                    self.diagnostics
-                        .add_unresolved_module(import_from, *level, module);
-                    Type::Unknown
-                }
-            }
+        let module_name = match module_name {
+            Ok(module_name) => module_name,
             Err(ModuleNameResolutionError::InvalidSyntax) => {
                 tracing::debug!("Failed to resolve import due to invalid syntax");
                 // Invalid syntax diagnostics are emitted elsewhere.
-                Type::Unknown
+                self.add_unknown_declaration_with_binding(alias.into(), definition);
+                return;
             }
             Err(ModuleNameResolutionError::TooManyDots) => {
                 tracing::debug!(
@@ -2319,7 +2278,8 @@ impl<'db> TypeInferenceBuilder<'db> {
                 );
                 self.diagnostics
                     .add_unresolved_module(import_from, *level, module);
-                Type::Unknown
+                self.add_unknown_declaration_with_binding(alias.into(), definition);
+                return;
             }
             Err(ModuleNameResolutionError::UnknownCurrentModule) => {
                 tracing::debug!(
@@ -2329,10 +2289,72 @@ impl<'db> TypeInferenceBuilder<'db> {
                 );
                 self.diagnostics
                     .add_unresolved_module(import_from, *level, module);
-                Type::Unknown
+                self.add_unknown_declaration_with_binding(alias.into(), definition);
+                return;
+            }
+        };
+
+        let Some(module_ty) = self.module_ty_from_name(&module_name) else {
+            self.diagnostics
+                .add_unresolved_module(import_from, *level, module);
+            self.add_unknown_declaration_with_binding(alias.into(), definition);
+            return;
+        };
+
+        let ast::Alias {
+            range: _,
+            name,
+            asname: _,
+        } = alias;
+
+        // Check if the symbol being imported is a submodule.  This won't get handled by the
+        // `Type::member` call below because it relies on the semantic index's `imported_modules`
+        // set.  The semantic index does not include information about `from...import` statements
+        // because there are two things it cannot determine while only inspecting the content of
+        // the current file:
+        //
+        //   - whether the imported symbol is an attribute or submodule
+        //   - whether the containing file is in a module or a package (needed to correctly resolve
+        //     relative imports)
+        //
+        // The first would be solvable by making it a _potentially_ imported modules set.  The
+        // second is not.
+        //
+        // Regardless, for now, we sidestep all of that by repeating the submodule-or-attribute
+        // check here when inferring types for a `from...import` statement.
+        if let Some(submodule_name) = ModuleName::new(name) {
+            let mut full_submodule_name = module_name.clone();
+            full_submodule_name.extend(&submodule_name);
+            if let Some(submodule_ty) = self.module_ty_from_name(&full_submodule_name) {
+                self.add_declaration_with_binding(
+                    alias.into(),
+                    definition,
+                    submodule_ty,
+                    submodule_ty,
+                );
+                return;
             }
+        }
+
+        // Otherwise load the requested attribute from the module.
+        let Symbol::Type(ty, boundness) = module_ty.member(self.db, name) else {
+            self.diagnostics.add_lint(
+                &UNRESOLVED_IMPORT,
+                AnyNodeRef::Alias(alias),
+                format_args!("Module `{module_name}` has no member `{name}`",),
+            );
+            self.add_unknown_declaration_with_binding(alias.into(), definition);
+            return;
         };
 
+        if boundness == Boundness::PossiblyUnbound {
+            self.diagnostics.add_lint(
+                &POSSIBLY_UNBOUND_IMPORT,
+                AnyNodeRef::Alias(alias),
+                format_args!("Member `{name}` of module `{module_name}` is possibly unbound",),
+            );
+        }
+
         self.add_declaration_with_binding(alias.into(), definition, ty, ty);
     }