Add dedicated struct for implicit imports (#5427)

## Summary

This was some feedback on a prior PR that I decided to act on
separately.
This commit is contained in:
Charlie Marsh 2023-06-28 14:55:43 -04:00 committed by GitHub
parent 0e12eb3071
commit 69c4b7fa11
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
10 changed files with 270 additions and 230 deletions

View file

@ -1,48 +1,32 @@
use std::collections::BTreeMap; use std::collections::BTreeMap;
use std::ffi::OsStr; use std::ffi::OsStr;
use std::fs; use std::io;
use std::path::{Path, PathBuf}; use std::path::{Path, PathBuf};
use crate::{native_module, py_typed}; use crate::{native_module, py_typed};
#[derive(Debug, Clone, PartialEq, Eq)] /// A map of the submodules that are present in a namespace package.
pub(crate) struct ImplicitImport { ///
/// Whether the implicit import is a stub file. /// Namespace packages lack an `__init__.py` file. So when resolving symbols from a namespace
pub(crate) is_stub_file: bool, /// 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<String, ImplicitImport>);
/// Whether the implicit import is a native module. impl ImplicitImports {
pub(crate) is_native_lib: bool, /// Find the "implicit" imports within the namespace package at the given path.
pub(crate) fn find(dir_path: &Path, exclusions: &[&Path]) -> io::Result<Self> {
/// The name of the implicit import (e.g., `os`). let mut submodules: BTreeMap<String, ImplicitImport> = BTreeMap::new();
pub(crate) name: String,
/// 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<py_typed::PyTypedInfo>,
}
/// Find the "implicit" imports within the namespace package at the given path.
pub(crate) fn find(dir_path: &Path, exclusions: &[&Path]) -> BTreeMap<String, ImplicitImport> {
let mut implicit_imports = BTreeMap::new();
// Enumerate all files and directories in the path, expanding links. // Enumerate all files and directories in the path, expanding links.
let Ok(entries) = fs::read_dir(dir_path) else { for entry in dir_path.read_dir()?.flatten() {
return implicit_imports; let file_type = entry.file_type()?;
};
for entry in entries.flatten() {
let path = entry.path(); let path = entry.path();
if exclusions.contains(&path.as_path()) { if exclusions.contains(&path.as_path()) {
continue; continue;
} }
let Ok(file_type) = entry.file_type() else {
continue;
};
// TODO(charlie): Support symlinks. // TODO(charlie): Support symlinks.
if file_type.is_file() { if file_type.is_file() {
// Add implicit file-based modules. // Add implicit file-based modules.
@ -68,20 +52,20 @@ pub(crate) fn find(dir_path: &Path, exclusions: &[&Path]) -> BTreeMap<String, Im
continue; continue;
}; };
let implicit_import = ImplicitImport {
is_stub_file: extension == "pyi",
is_native_lib,
name: name.to_string(),
path: path.clone(),
py_typed: None,
};
// Always prefer stub files over non-stub files. // Always prefer stub files over non-stub files.
if implicit_imports if submodules
.get(&implicit_import.name) .get(name)
.map_or(true, |implicit_import| !implicit_import.is_stub_file) .map_or(true, |implicit_import| !implicit_import.is_stub_file)
{ {
implicit_imports.insert(implicit_import.name.clone(), implicit_import); submodules.insert(
name.to_string(),
ImplicitImport {
is_stub_file: extension == "pyi",
is_native_lib,
path,
py_typed: None,
},
);
} }
} else if file_type.is_dir() { } else if file_type.is_dir() {
// Add implicit directory-based modules. // Add implicit directory-based modules.
@ -99,40 +83,94 @@ pub(crate) fn find(dir_path: &Path, exclusions: &[&Path]) -> BTreeMap<String, Im
let Some(name) = path.file_name().and_then(OsStr::to_str) else { let Some(name) = path.file_name().and_then(OsStr::to_str) else {
continue; continue;
}; };
submodules.insert(
let implicit_import = ImplicitImport { name.to_string(),
ImplicitImport {
is_stub_file, is_stub_file,
is_native_lib: false, is_native_lib: false,
name: name.to_string(),
path: path.clone(),
py_typed: py_typed::get_py_typed_info(&path), py_typed: py_typed::get_py_typed_info(&path),
}; path,
implicit_imports.insert(implicit_import.name.clone(), implicit_import); },
);
} }
} }
implicit_imports Ok(Self(submodules))
} }
/// Filter a map of implicit imports to only include those that were actually imported. /// Filter [`ImplicitImports`] to only those symbols that were imported.
pub(crate) fn filter( pub(crate) fn filter(&self, imported_symbols: &[String]) -> Option<Self> {
implicit_imports: &BTreeMap<String, ImplicitImport>, if self.is_empty() || imported_symbols.is_empty() {
imported_symbols: &[String],
) -> Option<BTreeMap<String, ImplicitImport>> {
if implicit_imports.is_empty() || imported_symbols.is_empty() {
return None; return None;
} }
let mut filtered_imports = BTreeMap::new(); let filtered: BTreeMap<String, ImplicitImport> = self
for implicit_import in implicit_imports.values() { .iter()
if imported_symbols.contains(&implicit_import.name) { .filter(|(name, _)| imported_symbols.contains(name))
filtered_imports.insert(implicit_import.name.clone(), implicit_import.clone()); .map(|(name, implicit_import)| (name.clone(), implicit_import.clone()))
} .collect();
}
if filtered_imports.len() == implicit_imports.len() { if filtered.len() == self.len() {
return None; return None;
} }
Some(filtered_imports) 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<Item = (&String, &ImplicitImport)> {
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<py_typed::PyTypedInfo>,
} }

View file

@ -1,9 +1,8 @@
//! Interface that describes the output of the import resolver. //! Interface that describes the output of the import resolver.
use std::collections::BTreeMap;
use std::path::PathBuf; use std::path::PathBuf;
use crate::implicit_imports::ImplicitImport; use crate::implicit_imports::ImplicitImports;
use crate::py_typed::PyTypedInfo; use crate::py_typed::PyTypedInfo;
#[derive(Debug, Clone, PartialEq, Eq)] #[derive(Debug, Clone, PartialEq, Eq)]
@ -69,11 +68,11 @@ pub(crate) struct ImportResult {
/// A map from file to resolved path, for all implicitly imported /// A map from file to resolved path, for all implicitly imported
/// modules that are part of a namespace package. /// modules that are part of a namespace package.
pub(crate) implicit_imports: BTreeMap<String, ImplicitImport>, pub(crate) implicit_imports: ImplicitImports,
/// Any implicit imports whose symbols were explicitly imported (i.e., via /// Any implicit imports whose symbols were explicitly imported (i.e., via
/// a `from x import y` statement). /// a `from x import y` statement).
pub(crate) filtered_implicit_imports: BTreeMap<String, ImplicitImport>, pub(crate) filtered_implicit_imports: ImplicitImports,
/// If the import resolved to a type hint (i.e., a `.pyi` file), then /// If the import resolved to a type hint (i.e., a `.pyi` file), then
/// a non-type-hint resolution will be stored here. /// a non-type-hint resolution will be stored here.
@ -105,8 +104,8 @@ impl ImportResult {
is_stdlib_typeshed_file: false, is_stdlib_typeshed_file: false,
is_third_party_typeshed_file: false, is_third_party_typeshed_file: false,
is_local_typings_file: false, is_local_typings_file: false,
implicit_imports: BTreeMap::default(), implicit_imports: ImplicitImports::default(),
filtered_implicit_imports: BTreeMap::default(), filtered_implicit_imports: ImplicitImports::default(),
non_stub_import_result: None, non_stub_import_result: None,
py_typed_info: None, py_typed_info: None,
package_directory: None, package_directory: None,

View file

@ -1,6 +1,5 @@
//! Resolves Python imports to their corresponding files on disk. //! Resolves Python imports to their corresponding files on disk.
use std::collections::BTreeMap;
use std::ffi::OsStr; use std::ffi::OsStr;
use std::path::{Path, PathBuf}; use std::path::{Path, PathBuf};
@ -8,10 +7,10 @@ use log::debug;
use crate::config::Config; use crate::config::Config;
use crate::execution_environment::ExecutionEnvironment; use crate::execution_environment::ExecutionEnvironment;
use crate::implicit_imports::ImplicitImport; use crate::implicit_imports::ImplicitImports;
use crate::import_result::{ImportResult, ImportType}; use crate::import_result::{ImportResult, ImportType};
use crate::module_descriptor::ImportModuleDescriptor; use crate::module_descriptor::ImportModuleDescriptor;
use crate::{host, implicit_imports, native_module, py_typed, search}; use crate::{host, native_module, py_typed, search};
#[allow(clippy::fn_params_excessive_bools)] #[allow(clippy::fn_params_excessive_bools)]
fn resolve_module_descriptor( fn resolve_module_descriptor(
@ -37,7 +36,7 @@ fn resolve_module_descriptor(
let mut is_stub_package = false; let mut is_stub_package = false;
let mut is_stub_file = false; let mut is_stub_file = false;
let mut is_native_lib = false; let mut is_native_lib = false;
let mut implicit_imports = BTreeMap::new(); let mut implicit_imports = None;
let mut package_directory = None; let mut package_directory = None;
let mut py_typed_info = None; let mut py_typed_info = None;
@ -60,7 +59,7 @@ fn resolve_module_descriptor(
is_namespace_package = true; is_namespace_package = true;
} }
implicit_imports = implicit_imports::find(&dir_path, &[&py_file_path, &pyi_file_path]); implicit_imports = ImplicitImports::find(&dir_path, &[&py_file_path, &pyi_file_path]).ok();
} else { } else {
for (i, part) in module_descriptor.name_parts.iter().enumerate() { for (i, part) in module_descriptor.name_parts.iter().enumerate() {
let is_first_part = i == 0; let is_first_part = i == 0;
@ -118,7 +117,8 @@ fn resolve_module_descriptor(
if is_init_file_present { if is_init_file_present {
implicit_imports = implicit_imports =
implicit_imports::find(&module_dir_path, &[&py_file_path, &pyi_file_path]); ImplicitImports::find(&module_dir_path, &[&py_file_path, &pyi_file_path])
.ok();
break; break;
} }
} }
@ -157,7 +157,7 @@ fn resolve_module_descriptor(
resolved_paths.push(PathBuf::new()); resolved_paths.push(PathBuf::new());
if is_last_part { if is_last_part {
implicit_imports = implicit_imports =
implicit_imports::find(&dir_path, &[&py_file_path, &pyi_file_path]); ImplicitImports::find(&dir_path, &[&py_file_path, &pyi_file_path]).ok();
is_namespace_package = true; is_namespace_package = true;
} }
} }
@ -191,8 +191,8 @@ fn resolve_module_descriptor(
is_stdlib_typeshed_file: false, is_stdlib_typeshed_file: false,
is_third_party_typeshed_file: false, is_third_party_typeshed_file: false,
is_local_typings_file: false, is_local_typings_file: false,
implicit_imports, implicit_imports: implicit_imports.unwrap_or_default(),
filtered_implicit_imports: BTreeMap::default(), filtered_implicit_imports: ImplicitImports::default(),
non_stub_import_result: None, non_stub_import_result: None,
py_typed_info, py_typed_info,
package_directory, package_directory,
@ -290,10 +290,10 @@ fn resolve_best_absolute_import<Host: host::Host>(
.as_os_str() .as_os_str()
.is_empty() .is_empty()
{ {
if is_namespace_package_resolved( if typings_import
module_descriptor, .implicit_imports
&typings_import.implicit_imports, .resolves_namespace_package(&module_descriptor.imported_symbols)
) { {
return Some(typings_import); return Some(typings_import);
} }
} else { } else {
@ -415,34 +415,6 @@ fn resolve_best_absolute_import<Host: host::Host>(
best_result_so_far best_result_so_far
} }
/// Determines whether a namespace package resolves all of the symbols
/// requested in the module descriptor. Namespace packages have no "__init__.py"
/// file, so the only way that symbols can be resolved is if submodules
/// are present. If specific symbols were requested, make sure they
/// are all satisfied by submodules (as listed in the implicit imports).
fn is_namespace_package_resolved(
module_descriptor: &ImportModuleDescriptor,
implicit_imports: &BTreeMap<String, ImplicitImport>,
) -> bool {
if !module_descriptor.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 &module_descriptor.imported_symbols {
if !implicit_imports.contains_key(symbol) {
return false;
}
}
} else if implicit_imports.is_empty() {
return false;
}
true
}
/// Finds the `typeshed` path for the given module descriptor. /// Finds the `typeshed` path for the given module descriptor.
/// ///
/// Supports both standard library and third-party `typeshed` lookups. /// Supports both standard library and third-party `typeshed` lookups.
@ -543,14 +515,14 @@ fn pick_best_import(
// imported symbols. // imported symbols.
if best_import_so_far.is_namespace_package && new_import.is_namespace_package { if best_import_so_far.is_namespace_package && new_import.is_namespace_package {
if !module_descriptor.imported_symbols.is_empty() { if !module_descriptor.imported_symbols.is_empty() {
if !is_namespace_package_resolved( if !best_import_so_far
module_descriptor, .implicit_imports
&best_import_so_far.implicit_imports, .resolves_namespace_package(&module_descriptor.imported_symbols)
) { {
if is_namespace_package_resolved( if new_import
module_descriptor, .implicit_imports
&new_import.implicit_imports, .resolves_namespace_package(&module_descriptor.imported_symbols)
) { {
return new_import; return new_import;
} }
@ -759,10 +731,10 @@ pub(crate) fn resolve_import<Host: host::Host>(
); );
if result.is_import_found { if result.is_import_found {
if let Some(implicit_imports) = implicit_imports::filter( if let Some(implicit_imports) = result
&result.implicit_imports, .implicit_imports
&module_descriptor.imported_symbols, .filter(&module_descriptor.imported_symbols)
) { {
result.implicit_imports = implicit_imports; result.implicit_imports = implicit_imports;
} }
return result; return result;

View file

@ -21,8 +21,12 @@ ImportResult {
is_stdlib_typeshed_file: false, is_stdlib_typeshed_file: false,
is_third_party_typeshed_file: false, is_third_party_typeshed_file: false,
is_local_typings_file: false, is_local_typings_file: false,
implicit_imports: {}, implicit_imports: ImplicitImports(
filtered_implicit_imports: {}, {},
),
filtered_implicit_imports: ImplicitImports(
{},
),
non_stub_import_result: None, non_stub_import_result: None,
py_typed_info: None, py_typed_info: None,
package_directory: None, package_directory: None,

View file

@ -23,8 +23,12 @@ ImportResult {
is_stdlib_typeshed_file: false, is_stdlib_typeshed_file: false,
is_third_party_typeshed_file: false, is_third_party_typeshed_file: false,
is_local_typings_file: false, is_local_typings_file: false,
implicit_imports: {}, implicit_imports: ImplicitImports(
filtered_implicit_imports: {}, {},
),
filtered_implicit_imports: ImplicitImports(
{},
),
non_stub_import_result: None, non_stub_import_result: None,
py_typed_info: None, py_typed_info: None,
package_directory: Some( package_directory: Some(

View file

@ -21,16 +21,19 @@ ImportResult {
is_stdlib_typeshed_file: false, is_stdlib_typeshed_file: false,
is_third_party_typeshed_file: false, is_third_party_typeshed_file: false,
is_local_typings_file: false, is_local_typings_file: false,
implicit_imports: { implicit_imports: ImplicitImports(
{
"orjson": ImplicitImport { "orjson": ImplicitImport {
is_stub_file: false, is_stub_file: false,
is_native_lib: true, is_native_lib: true,
name: "orjson",
path: "./resources/test/airflow/venv/lib/python3.11/site-packages/orjson/orjson.cpython-311-darwin.so", path: "./resources/test/airflow/venv/lib/python3.11/site-packages/orjson/orjson.cpython-311-darwin.so",
py_typed: None, py_typed: None,
}, },
}, },
filtered_implicit_imports: {}, ),
filtered_implicit_imports: ImplicitImports(
{},
),
non_stub_import_result: Some( non_stub_import_result: Some(
ImportResult { ImportResult {
is_relative: false, is_relative: false,
@ -51,16 +54,19 @@ ImportResult {
is_stdlib_typeshed_file: false, is_stdlib_typeshed_file: false,
is_third_party_typeshed_file: false, is_third_party_typeshed_file: false,
is_local_typings_file: false, is_local_typings_file: false,
implicit_imports: { implicit_imports: ImplicitImports(
{
"orjson": ImplicitImport { "orjson": ImplicitImport {
is_stub_file: false, is_stub_file: false,
is_native_lib: true, is_native_lib: true,
name: "orjson",
path: "./resources/test/airflow/venv/lib/python3.11/site-packages/orjson/orjson.cpython-311-darwin.so", path: "./resources/test/airflow/venv/lib/python3.11/site-packages/orjson/orjson.cpython-311-darwin.so",
py_typed: None, py_typed: None,
}, },
}, },
filtered_implicit_imports: {}, ),
filtered_implicit_imports: ImplicitImports(
{},
),
non_stub_import_result: None, non_stub_import_result: None,
py_typed_info: Some( py_typed_info: Some(
PyTypedInfo { PyTypedInfo {

View file

@ -26,8 +26,12 @@ ImportResult {
is_stdlib_typeshed_file: false, is_stdlib_typeshed_file: false,
is_third_party_typeshed_file: false, is_third_party_typeshed_file: false,
is_local_typings_file: false, is_local_typings_file: false,
implicit_imports: {}, implicit_imports: ImplicitImports(
filtered_implicit_imports: {}, {},
),
filtered_implicit_imports: ImplicitImports(
{},
),
non_stub_import_result: None, non_stub_import_result: None,
py_typed_info: None, py_typed_info: None,
package_directory: Some( package_directory: Some(

View file

@ -17,8 +17,12 @@ ImportResult {
is_stdlib_typeshed_file: false, is_stdlib_typeshed_file: false,
is_third_party_typeshed_file: false, is_third_party_typeshed_file: false,
is_local_typings_file: false, is_local_typings_file: false,
implicit_imports: {}, implicit_imports: ImplicitImports(
filtered_implicit_imports: {}, {},
),
filtered_implicit_imports: ImplicitImports(
{},
),
non_stub_import_result: None, non_stub_import_result: None,
py_typed_info: None, py_typed_info: None,
package_directory: None, package_directory: None,

View file

@ -23,8 +23,12 @@ ImportResult {
is_stdlib_typeshed_file: false, is_stdlib_typeshed_file: false,
is_third_party_typeshed_file: false, is_third_party_typeshed_file: false,
is_local_typings_file: false, is_local_typings_file: false,
implicit_imports: {}, implicit_imports: ImplicitImports(
filtered_implicit_imports: {}, {},
),
filtered_implicit_imports: ImplicitImports(
{},
),
non_stub_import_result: Some( non_stub_import_result: Some(
ImportResult { ImportResult {
is_relative: false, is_relative: false,
@ -47,8 +51,12 @@ ImportResult {
is_stdlib_typeshed_file: false, is_stdlib_typeshed_file: false,
is_third_party_typeshed_file: false, is_third_party_typeshed_file: false,
is_local_typings_file: false, is_local_typings_file: false,
implicit_imports: {}, implicit_imports: ImplicitImports(
filtered_implicit_imports: {}, {},
),
filtered_implicit_imports: ImplicitImports(
{},
),
non_stub_import_result: None, non_stub_import_result: None,
py_typed_info: None, py_typed_info: None,
package_directory: Some( package_directory: Some(

View file

@ -22,30 +22,31 @@ ImportResult {
is_stdlib_typeshed_file: false, is_stdlib_typeshed_file: false,
is_third_party_typeshed_file: false, is_third_party_typeshed_file: false,
is_local_typings_file: false, is_local_typings_file: false,
implicit_imports: { implicit_imports: ImplicitImports(
{
"base": ImplicitImport { "base": ImplicitImport {
is_stub_file: false, is_stub_file: false,
is_native_lib: false, is_native_lib: false,
name: "base",
path: "./resources/test/airflow/venv/lib/python3.11/site-packages/sqlalchemy/orm/base.py", path: "./resources/test/airflow/venv/lib/python3.11/site-packages/sqlalchemy/orm/base.py",
py_typed: None, py_typed: None,
}, },
"dependency": ImplicitImport { "dependency": ImplicitImport {
is_stub_file: false, is_stub_file: false,
is_native_lib: false, is_native_lib: false,
name: "dependency",
path: "./resources/test/airflow/venv/lib/python3.11/site-packages/sqlalchemy/orm/dependency.py", path: "./resources/test/airflow/venv/lib/python3.11/site-packages/sqlalchemy/orm/dependency.py",
py_typed: None, py_typed: None,
}, },
"query": ImplicitImport { "query": ImplicitImport {
is_stub_file: false, is_stub_file: false,
is_native_lib: false, is_native_lib: false,
name: "query",
path: "./resources/test/airflow/venv/lib/python3.11/site-packages/sqlalchemy/orm/query.py", path: "./resources/test/airflow/venv/lib/python3.11/site-packages/sqlalchemy/orm/query.py",
py_typed: None, py_typed: None,
}, },
}, },
filtered_implicit_imports: {}, ),
filtered_implicit_imports: ImplicitImports(
{},
),
non_stub_import_result: None, non_stub_import_result: None,
py_typed_info: None, py_typed_info: None,
package_directory: Some( package_directory: Some(