[ty] Add completions for submodule imports

While we did previously support submodule completions via our
`all_members` API, that only works when submodules are attributes of
their parent module. For example, `os.path`. But that didn't work when
the submodule was not an attribute of its parent. For example,
`http.client`. To make the latter work, we read the directory of the
parent module to discover its submodules.

crates/ty_ide/src/completion.rs

@@ -2396,6 +2396,48 @@ Cougar = 3
         test.assert_completions_include("Cheetah");
     }
 
+    #[test]
+    fn from_import_with_submodule1() {
+        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 from_import_with_vendored_submodule1() {
+        let test = cursor_test(
+            "\
+from http import <CURSOR>
+",
+        );
+        test.assert_completions_include("client");
+    }
+
+    #[test]
+    fn from_import_with_vendored_submodule2() {
+        let test = cursor_test(
+            "\
+from email import <CURSOR>
+",
+        );
+        test.assert_completions_include("mime");
+        test.assert_completions_do_not_include("base");
+    }
+
     #[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, get_size2::GetSize)]
@@ -85,6 +89,100 @@ 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).
+    ///
+    /// The names returned correspond to the "base" name of the module.
+    /// That is, `{self.name}.{basename}` should give the full module name.
+    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>> {
+        fn is_submodule(
+            is_dir: bool,
+            is_file: bool,
+            basename: Option<&str>,
+            extension: Option<&str>,
+        ) -> bool {
+            is_dir
+                || (is_file
+                    && matches!(extension, Some("py" | "pyi"))
+                    && !matches!(basename, Some("__init__.py" | "__init__.pyi")))
+        }
+
+        // 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")),
+            "expected package file `{:?}` to be `__init__.py` or `__init__.pyi`",
+            path.file_name(),
+        );
+
+        Some(match path.parent()? {
+            SystemOrVendoredPathRef::System(parent_directory) => db
+                .system()
+                .read_directory(parent_directory)
+                .inspect_err(|err| {
+                    tracing::debug!(
+                        "Failed to read {parent_directory:?} when looking for \
+                         its possible submodules: {err}"
+                    );
+                })
+                .ok()?
+                .flatten()
+                .filter(|entry| {
+                    let ty = entry.file_type();
+                    let path = entry.path();
+                    is_submodule(
+                        ty.is_directory(),
+                        ty.is_file(),
+                        path.file_name(),
+                        path.extension(),
+                    )
+                })
+                .filter_map(|entry| {
+                    let stem = entry.path().file_stem()?;
+                    is_identifier(stem).then(|| Name::from(stem))
+                })
+                .collect(),
+            SystemOrVendoredPathRef::Vendored(parent_directory) => db
+                .vendored()
+                .read_directory(parent_directory)
+                .into_iter()
+                .filter(|entry| {
+                    let ty = entry.file_type();
+                    let path = entry.path();
+                    is_submodule(
+                        ty.is_directory(),
+                        ty.is_file(),
+                        path.file_name(),
+                        path.extension(),
+                    )
+                })
+                .filter_map(|entry| {
+                    let stem = entry.path().file_stem()?;
+                    is_identifier(stem).then(|| Name::from(stem))
+                })
+                .collect(),
+        })
+    }
 }
 
 impl std::fmt::Debug for Module {

crates/ty_python_semantic/src/semantic_model.rs

@@ -69,14 +69,29 @@ impl<'db> SemanticModel<'db> {
         };
         let ty = Type::module_literal(self.db, self.file, &module);
         let builtin = module.is_known(KnownModule::Builtins);
-        crate::types::all_members(self.db, ty)
+
-            .into_iter()
+        let mut completions = vec![];
-            .map(|member| Completion {
+        for crate::types::Member { name, ty } in crate::types::all_members(self.db, ty) {
-                name: member.name,
+            completions.push(Completion { name, ty, builtin });
-                ty: member.ty,
+        }
+        for submodule_basename in module.all_submodules(self.db) {
+            let Some(basename) = ModuleName::new(submodule_basename.as_str()) else {
+                continue;
+            };
+            let mut submodule_name = module_name.clone();
+            submodule_name.extend(&basename);
+
+            let Some(submodule) = resolve_module(self.db, &submodule_name) else {
+                continue;
+            };
+            let ty = Type::module_literal(self.db, self.file, &submodule);
+            completions.push(Completion {
+                name: submodule_basename,
+                ty,
                 builtin,
-            })
+            });
-            .collect()
+        }
+        completions
     }
 
     /// Returns completions for symbols available in a `object.<CURSOR>` context.

crates/ty_python_semantic/src/types.rs

@@ -47,7 +47,7 @@ use crate::types::generics::{
     walk_partial_specialization, walk_specialization,
 };
 pub use crate::types::ide_support::{
-    CallSignatureDetails, all_members, call_signature_details, definition_kind_for_name,
+    CallSignatureDetails, Member, all_members, call_signature_details, definition_kind_for_name,
 };
 use crate::types::infer::infer_unpack_types;
 use crate::types::mro::{Mro, MroError, MroIterator};