Handle nested imports with as clause correctly

Author: dcreager

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

@@ -58,3 +58,36 @@ reveal_type(a.b.c.C)  # revealed: Literal[C]
 ```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

@@ -66,7 +66,5 @@ class C: ...
 
 ```py path=q.py
 import a as a
-import a.b
-# XXX: fix as
-b = a.b
+import a.b as b
 ```

crates/red_knot_python_semantic/src/types/infer.rs

@@ -2117,31 +2117,55 @@ impl<'db> TypeInferenceBuilder<'db> {
         let ast::Alias {
             range: _,
             name,
-            asname: _,
+            asname,
         } = alias;
 
-        // When importing a nested module a.b.c, we create a binding for (and need to infer the
-        // type of) the top-most parent module.  We also try to resolve the full nested module now,
-        // so that if we run into problems, any diagnostics are correctly associated with the
-        // import statement.
-        let module_ty = if let Some(module_name) = ModuleName::new(name) {
-            if self.module_ty_from_name(&module_name).is_none() {
-                self.diagnostics.add_unresolved_module(alias, 0, Some(name));
-            }
+        // 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,
+            );
+            return;
+        };
 
-            let bound_module_name =
-                ModuleName::new(module_name.components().next().unwrap()).unwrap();
-            if let Some(module) = self.module_ty_from_name(&bound_module_name) {
-                module
+        // 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,
+            );
+            return;
+        };
+
+        let binding_ty = if let Some(_) = asname {
+            // If we are renaming the imported module via an `as` clause, then we bind the resolved
+            // module's type to that name, even if that module is nested.
+            full_module_ty
+        } else if full_module_name.contains('.') {
+            // If there's no `as` clause and the imported module is nested, we're not going to bind
+            // the resolved module itself into the current scope; we're going to bind the top-most
+            // 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
             }
         } else {
-            tracing::debug!("Failed to resolve import due to invalid syntax");
-            Type::Unknown
+            // 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.
+            full_module_ty
         };
 
-        self.add_declaration_with_binding(alias.into(), definition, module_ty, module_ty);
+        self.add_declaration_with_binding(alias.into(), definition, binding_ty, binding_ty);
     }
 
     fn infer_import_from_statement(&mut self, import: &ast::StmtImportFrom) {