From a0ddf1f7c429a691414db2ec906af16d5c77d6b2 Mon Sep 17 00:00:00 2001
From: Andrew Gallant <andrew@astral.sh>
Date: Tue, 12 Aug 2025 14:10:02 -0400
Subject: [PATCH] [ty] Fix a bug when converting `ModulePath` to `ModuleName`

Previously, if the module was just `foo-stubs`, we'd skip over
stripping the `-stubs` suffix which would lead to us returning
`None`.

This function is now a little convoluted and could be simpler
if we did an intermediate allocation. But I kept the iterative
approach and added a special case to handle `foo-stubs`.
---
 .../src/module_resolver/path.rs               | 70 +++++++++++++++++--
 1 file changed, 66 insertions(+), 4 deletions(-)

diff --git a/crates/ty_python_semantic/src/module_resolver/path.rs b/crates/ty_python_semantic/src/module_resolver/path.rs
index 484b97bb69..c2c735bdd9 100644
--- a/crates/ty_python_semantic/src/module_resolver/path.rs
+++ b/crates/ty_python_semantic/src/module_resolver/path.rs
@@ -231,6 +231,10 @@ impl ModulePath {
 
     #[must_use]
     pub(crate) fn to_module_name(&self) -> Option<ModuleName> {
+        fn strip_stubs(component: &str) -> &str {
+            component.strip_suffix("-stubs").unwrap_or(component)
+        }
+
         let ModulePath {
             search_path: _,
             relative_path,
@@ -239,13 +243,28 @@ impl ModulePath {
             stdlib_path_to_module_name(relative_path)
         } else {
             let parent = relative_path.parent()?;
+            let name = relative_path.file_stem()?;
+            if parent.as_str().is_empty() {
+                // Stubs should only be stripped when there is no
+                // extension. e.g., `foo-stubs` should be stripped
+                // by not `foo-stubs.pyi`. In the latter case,
+                // `ModuleName::new` will fail (which is what we want).
+                return ModuleName::new(if relative_path.extension().is_some() {
+                    name
+                } else {
+                    strip_stubs(relative_path.as_str())
+                });
+            }
+
             let parent_components = parent.components().enumerate().map(|(index, component)| {
                 let component = component.as_str();
 
-                // For stub packages, strip the `-stubs` suffix from the first component
-                // because it isn't a valid module name part AND the module name is the name without the `-stubs`.
+                // For stub packages, strip the `-stubs` suffix from
+                // the first component because it isn't a valid module
+                // name part AND the module name is the name without
+                // the `-stubs`.
                 if index == 0 {
-                    component.strip_suffix("-stubs").unwrap_or(component)
+                    strip_stubs(component)
                 } else {
                     component
                 }
@@ -256,7 +275,7 @@ impl ModulePath {
             if skip_final_part {
                 ModuleName::from_components(parent_components)
             } else {
-                ModuleName::from_components(parent_components.chain(relative_path.file_stem()))
+                ModuleName::from_components(parent_components.chain([name]))
             }
         }
     }
@@ -1245,4 +1264,47 @@ mod tests {
         assert!(!xml_etree.is_directory(&resolver));
         assert!(!xml_etree.is_regular_package(&resolver));
     }
+
+    #[test]
+    fn strip_not_top_level_stubs_suffix() {
+        let TestCase { db, src, .. } = TestCaseBuilder::new().build();
+        let sp = SearchPath::first_party(db.system(), src).unwrap();
+        let mut mp = sp.to_module_path();
+        mp.push("foo-stubs");
+        mp.push("quux");
+        assert_eq!(
+            mp.to_module_name(),
+            Some(ModuleName::new_static("foo.quux").unwrap())
+        );
+    }
+
+    /// Tests that a module path of just `foo-stubs` will correctly be
+    /// converted to a module name of just `foo`.
+    ///
+    /// This is a regression test where this conversion ended up
+    /// treating the module path as invalid and returning `None` from
+    /// `ModulePath::to_module_name` instead.
+    #[test]
+    fn strip_top_level_stubs_suffix() {
+        let TestCase { db, src, .. } = TestCaseBuilder::new().build();
+        let sp = SearchPath::first_party(db.system(), src).unwrap();
+        let mut mp = sp.to_module_path();
+        mp.push("foo-stubs");
+        assert_eq!(
+            mp.to_module_name(),
+            Some(ModuleName::new_static("foo").unwrap())
+        );
+    }
+
+    /// Tests that paths like `foo-stubs.pyi` don't have `-stubs`
+    /// stripped. (And this leads to failing to create a `ModuleName`,
+    /// which is what we want.)
+    #[test]
+    fn no_strip_with_extension() {
+        let TestCase { db, src, .. } = TestCaseBuilder::new().build();
+        let sp = SearchPath::first_party(db.system(), src).unwrap();
+        let mut mp = sp.to_module_path();
+        mp.push("foo-stubs.pyi");
+        assert_eq!(mp.to_module_name(), None);
+    }
 }