From 5588c75d65e8edcedcead24eec192f4a3a8acba8 Mon Sep 17 00:00:00 2001
From: Micha Reiser <micha@reiser.io>
Date: Thu, 6 Feb 2025 14:08:20 +0000
Subject: [PATCH] [red-knot] Fix relative imports in `src.root` (#15990)

## Summary

Fixes https://github.com/astral-sh/ruff/issues/15989

Red Knot failed to resolve relative imports if the importing module is
located at a search path root.

The issue was that the module resolver returned an `Err(TooManyDots)` as
soon as the parent of the current module is `None` (which is the case
for a module at the search path root).
However, this is incorrect if a `tail` (a module name) exists.
---
 .../resources/mdtest/import/relative.md       | 30 +++++++++++++++++++
 .../src/types/infer.rs                        | 26 ++++++++++++----
 2 files changed, 50 insertions(+), 6 deletions(-)

diff --git a/crates/red_knot_python_semantic/resources/mdtest/import/relative.md b/crates/red_knot_python_semantic/resources/mdtest/import/relative.md
index cc34cbaf3e..f21c470646 100644
--- a/crates/red_knot_python_semantic/resources/mdtest/import/relative.md
+++ b/crates/red_knot_python_semantic/resources/mdtest/import/relative.md
@@ -218,3 +218,33 @@ import package
 # error: [unresolved-attribute] "Type `<module 'package'>` has no attribute `foo`"
 reveal_type(package.foo.X)  # revealed: Unknown
 ```
+
+## In the src-root
+
+`parser.py`:
+
+```py
+X: int = 42
+```
+
+`__main__.py`:
+
+```py
+from .parser import X
+
+reveal_type(X)  # revealed: int
+```
+
+## Beyond the src-root
+
+`parser.py`:
+
+```py
+X: int = 42
+```
+
+`__main__.py`:
+
+```py
+from ..parser import X  # error: [unresolved-import]
+```
diff --git a/crates/red_knot_python_semantic/src/types/infer.rs b/crates/red_knot_python_semantic/src/types/infer.rs
index da4793d7e7..8768c056c5 100644
--- a/crates/red_knot_python_semantic/src/types/infer.rs
+++ b/crates/red_knot_python_semantic/src/types/infer.rs
@@ -2513,18 +2513,32 @@ impl<'db> TypeInferenceBuilder<'db> {
             .ok_or(ModuleNameResolutionError::UnknownCurrentModule)?;
         let mut level = level.get();
         if module.kind().is_package() {
-            level -= 1;
+            level = level.saturating_sub(1);
         }
+
         let mut module_name = module.name().clone();
-        for _ in 0..level {
-            module_name = module_name
-                .parent()
-                .ok_or(ModuleNameResolutionError::TooManyDots)?;
+        let tail = tail
+            .map(|tail| ModuleName::new(tail).ok_or(ModuleNameResolutionError::InvalidSyntax))
+            .transpose()?;
+
+        for remaining_dots in (0..level).rev() {
+            if let Some(parent) = module_name.parent() {
+                module_name = parent;
+            } else if remaining_dots == 0 {
+                // If we reached a search path root and this was the last dot return the tail if any.
+                // If there's no tail, then we have a relative import that's too deep.
+                return tail.ok_or(ModuleNameResolutionError::TooManyDots);
+            } else {
+                // We're at a search path root. This is a too deep relative import if there's more than
+                // one dot remaining.
+                return Err(ModuleNameResolutionError::TooManyDots);
+            }
         }
+
         if let Some(tail) = tail {
-            let tail = ModuleName::new(tail).ok_or(ModuleNameResolutionError::InvalidSyntax)?;
             module_name.extend(&tail);
         }
+
         Ok(module_name)
     }