[red-knot] Add support for relative imports (#12910)

Co-authored-by: Carl Meyer <carl@astral.sh>

crates/red_knot_python_semantic/src/module_name.rs

@@ -168,6 +168,24 @@ impl ModuleName {
         };
         Some(Self(name))
     }
+
+    /// Extend `self` with the components of `other`
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use red_knot_python_semantic::ModuleName;
+    ///
+    /// let mut module_name = ModuleName::new_static("foo").unwrap();
+    /// module_name.extend(&ModuleName::new_static("bar").unwrap());
+    /// assert_eq!(&module_name, "foo.bar");
+    /// module_name.extend(&ModuleName::new_static("baz.eggs.ham").unwrap());
+    /// assert_eq!(&module_name, "foo.bar.baz.eggs.ham");
+    /// ```
+    pub fn extend(&mut self, other: &ModuleName) {
+        self.0.push('.');
+        self.0.push_str(other);
+    }
 }
 
 impl Deref for ModuleName {

crates/red_knot_python_semantic/src/module_resolver/mod.rs

@@ -2,7 +2,7 @@ use std::iter::FusedIterator;
 
 pub(crate) use module::Module;
 pub use resolver::resolve_module;
-pub(crate) use resolver::SearchPaths;
+pub(crate) use resolver::{file_to_module, SearchPaths};
 use ruff_db::system::SystemPath;
 pub use typeshed::vendored_typeshed_stubs;
 

crates/red_knot_python_semantic/src/module_resolver/module.rs

@@ -77,3 +77,9 @@ pub enum ModuleKind {
     /// A python package (`foo/__init__.py` or `foo/__init__.pyi`)
     Package,
 }
+
+impl ModuleKind {
+    pub const fn is_package(self) -> bool {
+        matches!(self, ModuleKind::Package)
+    }
+}

crates/red_knot_python_semantic/src/types/infer.rs

@@ -20,6 +20,8 @@
 //!
 //! Inferring types at any of the three region granularities returns a [`TypeInference`], which
 //! holds types for every [`Definition`] and expression within the inferred region.
+use std::num::NonZeroU32;
+
 use rustc_hash::FxHashMap;
 use salsa;
 use salsa::plumbing::AsId;
@@ -31,7 +33,7 @@ use ruff_python_ast::{ExprContext, TypeParams};
 
 use crate::builtins::builtins_scope;
 use crate::module_name::ModuleName;
-use crate::module_resolver::resolve_module;
+use crate::module_resolver::{file_to_module, resolve_module};
 use crate::semantic_index::ast_ids::{HasScopedAstId, HasScopedUseId, ScopedExpressionId};
 use crate::semantic_index::definition::{Definition, DefinitionKind, DefinitionNodeKey};
 use crate::semantic_index::expression::Expression;
@@ -822,7 +824,7 @@ impl<'db> TypeInferenceBuilder<'db> {
             asname: _,
         } = alias;
 
-        let module_ty = self.module_ty_from_name(name);
+        let module_ty = self.module_ty_from_name(ModuleName::new(name));
         self.types.definitions.insert(definition, module_ty);
     }
 
@@ -860,20 +862,68 @@ impl<'db> TypeInferenceBuilder<'db> {
         self.infer_optional_expression(cause.as_deref());
     }
 
+    /// Given a `from .foo import bar` relative import, resolve the relative module
+    /// we're importing `bar` from into an absolute [`ModuleName`]
+    /// using the name of the module we're currently analyzing.
+    ///
+    /// - `level` is the number of dots at the beginning of the relative module name:
+    ///   - `from .foo.bar import baz` => `level == 1`
+    ///   - `from ...foo.bar import baz` => `level == 3`
+    /// - `tail` is the relative module name stripped of all leading dots:
+    ///   - `from .foo import bar` => `tail == "foo"`
+    ///   - `from ..foo.bar import baz` => `tail == "foo.bar"`
+    fn relative_module_name(&self, tail: Option<&str>, level: NonZeroU32) -> Option<ModuleName> {
+        let Some(module) = file_to_module(self.db, self.file) else {
+            tracing::debug!("Failed to resolve file {:?} to a module", self.file);
+            return None;
+        };
+        let mut level = level.get();
+        if module.kind().is_package() {
+            level -= 1;
+        }
+        let mut module_name = module.name().to_owned();
+        for _ in 0..level {
+            module_name = module_name.parent()?;
+        }
+        if let Some(tail) = tail {
+            if let Some(valid_tail) = ModuleName::new(tail) {
+                module_name.extend(&valid_tail);
+            } else {
+                tracing::debug!("Failed to resolve relative import due to invalid syntax");
+                return None;
+            }
+        }
+        Some(module_name)
+    }
+
     fn infer_import_from_definition(
         &mut self,
         import_from: &ast::StmtImportFrom,
         alias: &ast::Alias,
         definition: Definition<'db>,
     ) {
-        let ast::StmtImportFrom { module, .. } = import_from;
-        let module_ty = if let Some(module) = module {
-            self.module_ty_from_name(module)
+        // 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;
+        tracing::trace!("Resolving imported object {alias:?} from statement {import_from:?}");
+        let module_name = if let Some(level) = NonZeroU32::new(*level) {
+            self.relative_module_name(module.as_deref(), level)
         } else {
-            // TODO support relative imports
-            Type::Unknown
+            let module_name = module
+                .as_ref()
+                .expect("Non-relative import should always have a non-None `module`!");
+            ModuleName::new(module_name)
         };
 
+        let module_ty = self.module_ty_from_name(module_name);
+
         let ast::Alias {
             range: _,
             name,
@@ -896,11 +946,10 @@ impl<'db> TypeInferenceBuilder<'db> {
         }
     }
 
-    fn module_ty_from_name(&self, name: &ast::Identifier) -> Type<'db> {
-        let module = ModuleName::new(&name.id).and_then(|name| resolve_module(self.db, name));
-        module
-            .map(|module| Type::Module(module.file()))
-            .unwrap_or(Type::Unbound)
+    fn module_ty_from_name(&self, module_name: Option<ModuleName>) -> Type<'db> {
+        module_name
+            .and_then(|module_name| resolve_module(self.db, module_name))
+            .map_or(Type::Unbound, |module| Type::Module(module.file()))
     }
 
     fn infer_decorator(&mut self, decorator: &ast::Decorator) -> Type<'db> {
@@ -1710,6 +1759,148 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn follow_relative_import_simple() -> anyhow::Result<()> {
+        let mut db = setup_db();
+
+        db.write_files([
+            ("src/package/__init__.py", ""),
+            ("src/package/foo.py", "X = 42"),
+            ("src/package/bar.py", "from .foo import X"),
+        ])?;
+
+        assert_public_ty(&db, "src/package/bar.py", "X", "Literal[42]");
+
+        Ok(())
+    }
+
+    #[test]
+    fn follow_nonexistent_relative_import_simple() -> anyhow::Result<()> {
+        let mut db = setup_db();
+
+        db.write_files([
+            ("src/package/__init__.py", ""),
+            ("src/package/bar.py", "from .foo import X"),
+        ])?;
+
+        assert_public_ty(&db, "src/package/bar.py", "X", "Unbound");
+
+        Ok(())
+    }
+
+    #[test]
+    fn follow_relative_import_dotted() -> anyhow::Result<()> {
+        let mut db = setup_db();
+
+        db.write_files([
+            ("src/package/__init__.py", ""),
+            ("src/package/foo/bar/baz.py", "X = 42"),
+            ("src/package/bar.py", "from .foo.bar.baz import X"),
+        ])?;
+
+        assert_public_ty(&db, "src/package/bar.py", "X", "Literal[42]");
+
+        Ok(())
+    }
+
+    #[test]
+    fn follow_relative_import_bare_to_package() -> anyhow::Result<()> {
+        let mut db = setup_db();
+
+        db.write_files([
+            ("src/package/__init__.py", "X = 42"),
+            ("src/package/bar.py", "from . import X"),
+        ])?;
+
+        assert_public_ty(&db, "src/package/bar.py", "X", "Literal[42]");
+
+        Ok(())
+    }
+
+    #[test]
+    fn follow_nonexistent_relative_import_bare_to_package() -> anyhow::Result<()> {
+        let mut db = setup_db();
+        db.write_files([("src/package/bar.py", "from . import X")])?;
+        assert_public_ty(&db, "src/package/bar.py", "X", "Unbound");
+        Ok(())
+    }
+
+    #[ignore = "TODO: Submodule imports possibly not supported right now?"]
+    #[test]
+    fn follow_relative_import_bare_to_module() -> anyhow::Result<()> {
+        let mut db = setup_db();
+
+        db.write_files([
+            ("src/package/__init__.py", ""),
+            ("src/package/foo.py", "X = 42"),
+            ("src/package/bar.py", "from . import foo; y = foo.X"),
+        ])?;
+
+        assert_public_ty(&db, "src/package/bar.py", "y", "Literal[42]");
+
+        Ok(())
+    }
+
+    #[ignore = "TODO: Submodule imports possibly not supported right now?"]
+    #[test]
+    fn follow_nonexistent_import_bare_to_module() -> anyhow::Result<()> {
+        let mut db = setup_db();
+
+        db.write_files([
+            ("src/package/__init__.py", ""),
+            ("src/package/bar.py", "from . import foo"),
+        ])?;
+
+        assert_public_ty(&db, "src/package/bar.py", "foo", "Unbound");
+
+        Ok(())
+    }
+
+    #[test]
+    fn follow_relative_import_from_dunder_init() -> anyhow::Result<()> {
+        let mut db = setup_db();
+
+        db.write_files([
+            ("src/package/__init__.py", "from .foo import X"),
+            ("src/package/foo.py", "X = 42"),
+        ])?;
+
+        assert_public_ty(&db, "src/package/__init__.py", "X", "Literal[42]");
+
+        Ok(())
+    }
+
+    #[test]
+    fn follow_nonexistent_relative_import_from_dunder_init() -> anyhow::Result<()> {
+        let mut db = setup_db();
+        db.write_files([("src/package/__init__.py", "from .foo import X")])?;
+        assert_public_ty(&db, "src/package/__init__.py", "X", "Unbound");
+        Ok(())
+    }
+
+    #[test]
+    fn follow_very_relative_import() -> anyhow::Result<()> {
+        let mut db = setup_db();
+
+        db.write_files([
+            ("src/package/__init__.py", ""),
+            ("src/package/foo.py", "X = 42"),
+            (
+                "src/package/subpackage/subsubpackage/bar.py",
+                "from ...foo import X",
+            ),
+        ])?;
+
+        assert_public_ty(
+            &db,
+            "src/package/subpackage/subsubpackage/bar.py",
+            "X",
+            "Literal[42]",
+        );
+
+        Ok(())
+    }
+
     #[test]
     fn resolve_base_class_by_name() -> anyhow::Result<()> {
         let mut db = setup_db();

crates/ruff_benchmark/benches/red_knot.rs

@@ -89,7 +89,7 @@ fn benchmark_incremental(criterion: &mut Criterion) {
                 let Case { db, parser, .. } = case;
                 let result = db.check_file(*parser).unwrap();
 
-                assert_eq!(result.len(), 111);
+                assert_eq!(result.len(), 29);
             },
             BatchSize::SmallInput,
         );
@@ -104,7 +104,7 @@ fn benchmark_cold(criterion: &mut Criterion) {
                 let Case { db, parser, .. } = case;
                 let result = db.check_file(*parser).unwrap();
 
-                assert_eq!(result.len(), 111);
+                assert_eq!(result.len(), 29);
             },
             BatchSize::SmallInput,
         );