use std::collections::BTreeMap; use std::ffi::OsStr; use std::io; use std::path::{Path, PathBuf}; use crate::{native_module, py_typed}; /// A map of the submodules that are present in a namespace package. /// /// Namespace packages lack an `__init__.py` file. So when resolving symbols from a namespace /// package, the symbols must be present as submodules. This map contains the submodules that are /// present in the namespace package, keyed by their module name. #[derive(Debug, Clone, PartialEq, Eq, Default)] pub(crate) struct ImplicitImports(BTreeMap); impl ImplicitImports { /// Find the "implicit" imports within the namespace package at the given path. pub(crate) fn find(dir_path: &Path, exclusions: &[&Path]) -> io::Result { let mut submodules: BTreeMap = BTreeMap::new(); // Enumerate all files and directories in the path, expanding links. for entry in dir_path.read_dir()?.flatten() { let file_type = entry.file_type()?; let path = entry.path(); if exclusions.contains(&path.as_path()) { continue; } // TODO(charlie): Support symlinks. if file_type.is_file() { // Add implicit file-based modules. let Some(extension) = path.extension() else { continue; }; let (file_stem, is_native_lib) = if extension == "py" || extension == "pyi" { // E.g., `foo.py` becomes `foo`. let file_stem = path.file_stem().and_then(OsStr::to_str); let is_native_lib = false; (file_stem, is_native_lib) } else if native_module::is_native_module_file_extension(extension) { // E.g., `foo.abi3.so` becomes `foo`. let file_stem = native_module::native_module_name(&path); let is_native_lib = true; (file_stem, is_native_lib) } else { continue; }; let Some(name) = file_stem else { continue; }; // Always prefer stub files over non-stub files. if submodules .get(name) .map_or(true, |implicit_import| !implicit_import.is_stub_file) { submodules.insert( name.to_string(), ImplicitImport { is_stub_file: extension == "pyi", is_native_lib, path, py_typed: None, }, ); } } else if file_type.is_dir() { // Add implicit directory-based modules. let py_file_path = path.join("__init__.py"); let pyi_file_path = path.join("__init__.pyi"); let (path, is_stub_file) = if py_file_path.exists() { (py_file_path, false) } else if pyi_file_path.exists() { (pyi_file_path, true) } else { continue; }; let Some(name) = path.file_name().and_then(OsStr::to_str) else { continue; }; submodules.insert( name.to_string(), ImplicitImport { is_stub_file, is_native_lib: false, py_typed: py_typed::get_py_typed_info(&path), path, }, ); } } Ok(Self(submodules)) } /// Filter [`ImplicitImports`] to only those symbols that were imported. pub(crate) fn filter(&self, imported_symbols: &[String]) -> Option { if self.is_empty() || imported_symbols.is_empty() { return None; } let filtered: BTreeMap = self .iter() .filter(|(name, _)| imported_symbols.contains(name)) .map(|(name, implicit_import)| (name.clone(), implicit_import.clone())) .collect(); if filtered.len() == self.len() { return None; } Some(Self(filtered)) } /// Returns `true` if the [`ImplicitImports`] resolves all the symbols requested by a /// module descriptor. pub(crate) fn resolves_namespace_package(&self, imported_symbols: &[String]) -> bool { if !imported_symbols.is_empty() { // TODO(charlie): Pyright uses: // // ```typescript // !Array.from(moduleDescriptor.importedSymbols.keys()).some((symbol) => implicitImports.has(symbol))` // ``` // // However, that only checks if _any_ of the symbols are in the implicit imports. for symbol in imported_symbols { if !self.has(symbol) { return false; } } } else if self.is_empty() { return false; } true } /// Returns `true` if the module is present in the namespace package. pub(crate) fn has(&self, name: &str) -> bool { self.0.contains_key(name) } /// Returns the number of implicit imports in the namespace package. pub(crate) fn len(&self) -> usize { self.0.len() } /// Returns `true` if there are no implicit imports in the namespace package. pub(crate) fn is_empty(&self) -> bool { self.0.is_empty() } /// Returns an iterator over the implicit imports in the namespace package. pub(crate) fn iter(&self) -> impl Iterator { self.0.iter() } } #[derive(Debug, Clone, PartialEq, Eq)] pub(crate) struct ImplicitImport { /// Whether the implicit import is a stub file. pub(crate) is_stub_file: bool, /// Whether the implicit import is a native module. pub(crate) is_native_lib: bool, /// The path to the implicit import. pub(crate) path: PathBuf, /// The `py.typed` information for the implicit import, if any. pub(crate) py_typed: Option, }