diff --git a/crates/ty_ide/src/completion.rs b/crates/ty_ide/src/completion.rs index 312a715d8f..5988b6a930 100644 --- a/crates/ty_ide/src/completion.rs +++ b/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 ") + .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 +", + ); + test.assert_completions_include("client"); + } + + #[test] + fn from_import_with_vendored_submodule2() { + let test = cursor_test( + "\ +from email import +", + ); + test.assert_completions_include("mime"); + test.assert_completions_do_not_include("base"); + } + #[test] fn import_submodule_not_attribute1() { let test = cursor_test( diff --git a/crates/ty_python_semantic/src/module_resolver/module.rs b/crates/ty_python_semantic/src/module_resolver/module.rs index 9241095703..c1d322f90e 100644 --- a/crates/ty_python_semantic/src/module_resolver/module.rs +++ b/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 { + self.all_submodules_inner(db).unwrap_or_default() + } + + fn all_submodules_inner(&self, db: &dyn Db) -> Option> { + 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 { diff --git a/crates/ty_python_semantic/src/semantic_model.rs b/crates/ty_python_semantic/src/semantic_model.rs index e26992a39a..3f2ef429e3 100644 --- a/crates/ty_python_semantic/src/semantic_model.rs +++ b/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() - .map(|member| Completion { - name: member.name, - ty: member.ty, + + let mut completions = vec![]; + for crate::types::Member { name, ty } in crate::types::all_members(self.db, ty) { + completions.push(Completion { name, ty, builtin }); + } + 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.` context. diff --git a/crates/ty_python_semantic/src/types.rs b/crates/ty_python_semantic/src/types.rs index 3463ee933e..8a1c6662b0 100644 --- a/crates/ty_python_semantic/src/types.rs +++ b/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};