Very rough initial version of completions for submodule imports

crates/ruff_db/src/vendored/path.rs

@@ -17,6 +17,10 @@ impl VendoredPath {
         unsafe { &*(path as *const Utf8Path as *const VendoredPath) }
     }
 
+    pub fn file_name(&self) -> Option<&str> {
+        self.0.file_name()
+    }
+
     pub fn to_path_buf(&self) -> VendoredPathBuf {
         VendoredPathBuf(self.0.to_path_buf())
     }

crates/ty_ide/src/completion.rs

@@ -2180,6 +2180,27 @@ Cougar = 3
         test.assert_completions_include("Cheetah");
     }
 
+    #[test]
+    fn from_import_with_submodule() {
+        let test = CursorTest::builder()
+            .source("main.py", "from package import <CURSOR>")
+            .source("package/__init__.py", "")
+            .source("package/foo.py", "")
+            .source("package/bar.pyi", "")
+            .source("package/foo-bar.py", "")
+            .source("package/data.txt", "")
+            .source("package/sub/__init__.py", "")
+            .source("package/not-a-submodule/__init__.py", "")
+            .build();
+
+        test.assert_completions_include("foo");
+        test.assert_completions_include("bar");
+        test.assert_completions_include("sub");
+        test.assert_completions_do_not_include("foo-bar");
+        test.assert_completions_do_not_include("data");
+        test.assert_completions_do_not_include("not-a-submodule");
+    }
+
     #[test]
     fn import_submodule_not_attribute1() {
         let test = cursor_test(

crates/ty_python_semantic/src/module_resolver/module.rs

@@ -3,9 +3,13 @@ use std::str::FromStr;
 use std::sync::Arc;
 
 use ruff_db::files::File;
+use ruff_python_ast::name::Name;
+use ruff_python_stdlib::identifiers::is_identifier;
 
 use super::path::SearchPath;
+use crate::Db;
 use crate::module_name::ModuleName;
+use crate::module_resolver::path::SystemOrVendoredPathRef;
 
 /// Representation of a Python module.
 #[derive(Clone, PartialEq, Eq, Hash)]
@@ -85,6 +89,66 @@ impl Module {
             ModuleInner::NamespacePackage { .. } => ModuleKind::Package,
         }
     }
+
+    /// Return a list of all submodules of this module.
+    ///
+    /// Returns an empty list if the module is not a package, if it is an empty package,
+    /// or if it is a namespace package (one without an `__init__.py` or `__init__.pyi` file).
+    pub fn all_submodules(&self, db: &dyn Db) -> Vec<Name> {
+        self.all_submodules_inner(db).unwrap_or_default()
+    }
+
+    fn all_submodules_inner(&self, db: &dyn Db) -> Option<Vec<Name>> {
+        // It would be complex and expensive to compute all submodules for
+        // namespace packages, since a namespace package doesn't correspond
+        // to a single file; it can span multiple directories across multiple
+        // search paths. For now, we only compute submodules for traditional
+        // packages that exist in a single directory on a single search path.
+        let ModuleInner::FileModule {
+            kind: ModuleKind::Package,
+            file,
+            ..
+        } = &*self.inner
+        else {
+            return None;
+        };
+
+        let path = SystemOrVendoredPathRef::try_from_file(db, *file)?;
+
+        debug_assert!(
+            matches!(path.file_name(), Some("__init__.py" | "__init__.pyi"),),
+            "{:?}",
+            path.file_name()
+        );
+
+        let parent_directory = path.parent()?;
+
+        // TODO: it should be possible to compute submodules for vendored packages as well.
+        let SystemOrVendoredPathRef::System(parent_directory) = parent_directory else {
+            return None;
+        };
+
+        let directory_iterator = db.system().read_directory(parent_directory).ok()?;
+
+        let submodules = directory_iterator
+            .flatten()
+            .filter(|entry| {
+                entry.file_type().is_directory()
+                    || (entry.file_type().is_file()
+                        && matches!(entry.path().extension(), Some("py" | "pyi"))
+                        && !matches!(
+                            entry.path().file_name(),
+                            Some("__init__.py" | "__init__.pyi")
+                        ))
+            })
+            .filter_map(|entry| {
+                let stem = entry.path().file_stem()?;
+                is_identifier(stem).then(|| Name::from(stem))
+            })
+            .collect();
+
+        Some(submodules)
+    }
 }
 
 impl std::fmt::Debug for Module {

crates/ty_python_semantic/src/module_resolver/path.rs

@@ -4,7 +4,7 @@ use std::fmt;
 use std::sync::Arc;
 
 use camino::{Utf8Path, Utf8PathBuf};
-use ruff_db::files::{File, FileError, system_path_to_file, vendored_path_to_file};
+use ruff_db::files::{File, FileError, FilePath, system_path_to_file, vendored_path_to_file};
 use ruff_db::system::{System, SystemPath, SystemPathBuf};
 use ruff_db::vendored::{VendoredPath, VendoredPathBuf};
 
@@ -673,6 +673,49 @@ impl fmt::Display for SearchPath {
     }
 }
 
+#[derive(Debug, Clone)]
+pub(super) enum SystemOrVendoredPathRef<'db> {
+    System(&'db SystemPath),
+    Vendored(&'db VendoredPath),
+}
+
+impl<'db> SystemOrVendoredPathRef<'db> {
+    pub(super) fn try_from_file(db: &'db dyn Db, file: File) -> Option<Self> {
+        match file.path(db.upcast()) {
+            FilePath::System(system) => Some(Self::System(system)),
+            FilePath::Vendored(vendored) => Some(Self::Vendored(vendored)),
+            FilePath::SystemVirtual(_) => None,
+        }
+    }
+
+    #[cfg(debug_assertions)]
+    pub(super) fn file_name(&self) -> Option<&str> {
+        match self {
+            Self::System(system) => system.file_name(),
+            Self::Vendored(vendored) => vendored.file_name(),
+        }
+    }
+
+    pub(super) fn parent<'a>(&'a self) -> Option<SystemOrVendoredPathRef<'a>>
+    where
+        'a: 'db,
+    {
+        match self {
+            Self::System(system) => system.parent().map(Self::System),
+            Self::Vendored(vendored) => vendored.parent().map(Self::Vendored),
+        }
+    }
+}
+
+impl std::fmt::Display for SystemOrVendoredPathRef<'_> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            SystemOrVendoredPathRef::System(system) => system.fmt(f),
+            SystemOrVendoredPathRef::Vendored(vendored) => vendored.fmt(f),
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use ruff_db::Db;

crates/ty_python_semantic/src/module_resolver/resolver.rs

@@ -9,7 +9,7 @@ use rustc_hash::{FxBuildHasher, FxHashSet};
 
 use ruff_db::files::{File, FilePath, FileRootKind};
 use ruff_db::system::{DirectoryEntry, System, SystemPath, SystemPathBuf};
-use ruff_db::vendored::{VendoredFileSystem, VendoredPath};
+use ruff_db::vendored::VendoredFileSystem;
 use ruff_python_ast::PythonVersion;
 
 use crate::db::Db;
@@ -21,7 +21,7 @@ use crate::{
 };
 
 use super::module::{Module, ModuleKind};
-use super::path::{ModulePath, SearchPath, SearchPathValidationError};
+use super::path::{ModulePath, SearchPath, SearchPathValidationError, SystemOrVendoredPathRef};
 
 /// Resolves a module name to a module.
 pub fn resolve_module(db: &dyn Db, module_name: &ModuleName) -> Option<Module> {
@@ -81,33 +81,13 @@ pub(crate) fn path_to_module(db: &dyn Db, path: &FilePath) -> Option<Module> {
     file_to_module(db, file)
 }
 
-#[derive(Debug, Clone, Copy)]
-enum SystemOrVendoredPathRef<'a> {
-    System(&'a SystemPath),
-    Vendored(&'a VendoredPath),
-}
-
-impl std::fmt::Display for SystemOrVendoredPathRef<'_> {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        match self {
-            SystemOrVendoredPathRef::System(system) => system.fmt(f),
-            SystemOrVendoredPathRef::Vendored(vendored) => vendored.fmt(f),
-        }
-    }
-}
-
 /// Resolves the module for the file with the given id.
 ///
 /// Returns `None` if the file is not a module locatable via any of the known search paths.
 #[salsa::tracked]
 pub(crate) fn file_to_module(db: &dyn Db, file: File) -> Option<Module> {
     let _span = tracing::trace_span!("file_to_module", ?file).entered();
-
-    let path = match file.path(db.upcast()) {
-        FilePath::System(system) => SystemOrVendoredPathRef::System(system),
-        FilePath::Vendored(vendored) => SystemOrVendoredPathRef::Vendored(vendored),
-        FilePath::SystemVirtual(_) => return None,
-    };
+    let path = SystemOrVendoredPathRef::try_from_file(db, file)?;
 
     let module_name = search_paths(db).find_map(|candidate| {
         let relative_path = match path {

crates/ty_python_semantic/src/semantic_model.rs

@@ -63,7 +63,10 @@ impl<'db> SemanticModel<'db> {
             return vec![];
         };
         let ty = Type::module_literal(self.db, self.file, &module);
-        crate::types::all_members(self.db, ty).into_iter().collect()
+        crate::types::all_members(self.db, ty)
+            .into_iter()
+            .chain(module.all_submodules(self.db))
+            .collect()
     }
 
     /// Returns completions for symbols available in a `object.<CURSOR>` context.