diff --git a/Cargo.lock b/Cargo.lock index 287b62671a..57c612d87e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -674,6 +674,19 @@ version = "0.3.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a357d28ed41a50f9c765dbfe56cbc04a64e53e5fc58ba79fbc34c10ef3df831f" +[[package]] +name = "env_logger" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85cdab6a89accf66733ad5a1693a4dcced6aeff64602b634530dd73c1f3ee9f0" +dependencies = [ + "humantime", + "is-terminal", + "log", + "regex", + "termcolor", +] + [[package]] name = "equivalent" version = "1.0.0" @@ -874,6 +887,12 @@ version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "dfa686283ad6dd069f105e5ab091b04c62850d3e4cf5d67debad1933f55023df" +[[package]] +name = "humantime" +version = "2.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a3a5bfb195931eeb336b2a7b4d761daec841b97f947d34394601737a7bba5e4" + [[package]] name = "iana-time-zone" version = "0.1.57" @@ -2071,6 +2090,15 @@ dependencies = [ "similar", ] +[[package]] +name = "ruff_python_resolver" +version = "0.0.0" +dependencies = [ + "env_logger", + "log", + "tempfile", +] + [[package]] name = "ruff_python_semantic" version = "0.0.0" @@ -2539,6 +2567,15 @@ dependencies = [ "windows-sys 0.48.0", ] +[[package]] +name = "termcolor" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "be55cf8942feac5c765c2c993422806843c9a9a45d4d5c407ad6dd2ea95eb9b6" +dependencies = [ + "winapi-util", +] + [[package]] name = "terminfo" version = "0.8.0" diff --git a/LICENSE b/LICENSE index 68a9d4958c..8ffd09c5f2 100644 --- a/LICENSE +++ b/LICENSE @@ -1224,6 +1224,32 @@ are: SOFTWARE. """ +- Pyright, licensed as follows: + """ + MIT License + + Pyright - A static type checker for the Python language + Copyright (c) Microsoft Corporation. All rights reserved. + + Permission is hereby granted, free of charge, to any person obtaining a copy + of this software and associated documentation files (the "Software"), to deal + in the Software without restriction, including without limitation the rights + to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + copies of the Software, and to permit persons to whom the Software is + furnished to do so, subject to the following conditions: + + The above copyright notice and this permission notice shall be included in all + copies or substantial portions of the Software. + + THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE + SOFTWARE + """ + - rust-analyzer/text-size, licensed under the MIT license: """ Permission is hereby granted, free of charge, to any diff --git a/README.md b/README.md index 5c77173aad..46749a3e84 100644 --- a/README.md +++ b/README.md @@ -330,9 +330,11 @@ We're grateful to the maintainers of these tools for their work, and for all the value they've provided to the Python community. Ruff's autoformatter is built on a fork of Rome's [`rome_formatter`](https://github.com/rome/tools/tree/main/crates/rome_formatter), -and again draws on both the APIs and implementation details of [Rome](https://github.com/rome/tools), +and again draws on both API and implementation details from [Rome](https://github.com/rome/tools), [Prettier](https://github.com/prettier/prettier), and [Black](https://github.com/psf/black). +Ruff's import resolver is based on the import resolution algorithm from [Pyright](https://github.com/microsoft/pyright). + Ruff is also influenced by a number of tools outside the Python ecosystem, like [Clippy](https://github.com/rust-lang/rust-clippy) and [ESLint](https://github.com/eslint/eslint). diff --git a/crates/ruff_python_resolver/Cargo.toml b/crates/ruff_python_resolver/Cargo.toml new file mode 100644 index 0000000000..2d454e88cb --- /dev/null +++ b/crates/ruff_python_resolver/Cargo.toml @@ -0,0 +1,21 @@ +[package] +name = "ruff_python_resolver" +version = "0.0.0" +description = "A Python module resolver for Ruff" +publish = false +authors = { workspace = true } +edition = { workspace = true } +rust-version = { workspace = true } +homepage = { workspace = true } +documentation = { workspace = true } +repository = { workspace = true } +license = { workspace = true } + +[lib] + +[dependencies] +log = { workspace = true } + +[dev-dependencies] +env_logger = "0.10.0" +tempfile = "3.6.0" diff --git a/crates/ruff_python_resolver/src/config.rs b/crates/ruff_python_resolver/src/config.rs new file mode 100644 index 0000000000..0ae2790683 --- /dev/null +++ b/crates/ruff_python_resolver/src/config.rs @@ -0,0 +1,26 @@ +use std::path::PathBuf; + +use crate::python_version::PythonVersion; + +pub(crate) struct Config { + /// Path to python interpreter. + pub(crate) python_path: Option, + + /// Path to use for typeshed definitions. + pub(crate) typeshed_path: Option, + + /// Path to custom typings (stub) modules. + pub(crate) stub_path: Option, + + /// Path to a directory containing one or more virtual environment + /// directories. This is used in conjunction with the "venv" name in + /// the config file to identify the python environment used for resolving + /// third-party modules. + pub(crate) venv_path: Option, + + /// Default venv environment. + pub(crate) venv: Option, + + /// Default Python version. Can be overridden by ExecutionEnvironment. + pub(crate) default_python_version: Option, +} diff --git a/crates/ruff_python_resolver/src/execution_environment.rs b/crates/ruff_python_resolver/src/execution_environment.rs new file mode 100644 index 0000000000..b969ddc42b --- /dev/null +++ b/crates/ruff_python_resolver/src/execution_environment.rs @@ -0,0 +1,19 @@ +use std::path::PathBuf; + +use crate::python_platform::PythonPlatform; +use crate::python_version::PythonVersion; + +#[derive(Debug)] +pub(crate) struct ExecutionEnvironment { + /// The root directory of the execution environment. + pub(crate) root: PathBuf, + + /// The Python version of the execution environment. + pub(crate) python_version: PythonVersion, + + /// The Python platform of the execution environment. + pub(crate) python_platform: PythonPlatform, + + /// The extra search paths of the execution environment. + pub(crate) extra_paths: Vec, +} diff --git a/crates/ruff_python_resolver/src/host.rs b/crates/ruff_python_resolver/src/host.rs new file mode 100644 index 0000000000..be9b0a5e60 --- /dev/null +++ b/crates/ruff_python_resolver/src/host.rs @@ -0,0 +1,43 @@ +//! Expose the host environment to the resolver. + +use std::path::PathBuf; + +use crate::python_platform::PythonPlatform; +use crate::python_version::PythonVersion; + +/// A trait to expose the host environment to the resolver. +pub(crate) trait Host { + /// The search paths to use when resolving Python modules. + fn python_search_paths(&self) -> Vec; + + /// The Python version to use when resolving Python modules. + fn python_version(&self) -> PythonVersion; + + /// The OS platform to use when resolving Python modules. + fn python_platform(&self) -> PythonPlatform; +} + +/// A host that exposes a fixed set of search paths. +pub(crate) struct StaticHost { + search_paths: Vec, +} + +impl StaticHost { + pub(crate) fn new(search_paths: Vec) -> Self { + Self { search_paths } + } +} + +impl Host for StaticHost { + fn python_search_paths(&self) -> Vec { + self.search_paths.clone() + } + + fn python_version(&self) -> PythonVersion { + PythonVersion::Py312 + } + + fn python_platform(&self) -> PythonPlatform { + PythonPlatform::Darwin + } +} diff --git a/crates/ruff_python_resolver/src/implicit_imports.rs b/crates/ruff_python_resolver/src/implicit_imports.rs new file mode 100644 index 0000000000..2c42f5450a --- /dev/null +++ b/crates/ruff_python_resolver/src/implicit_imports.rs @@ -0,0 +1,150 @@ +use std::collections::HashMap; +use std::ffi::OsStr; +use std::fs; +use std::path::{Path, PathBuf}; + +use crate::{native_module, py_typed}; + +#[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 name of the implicit import (e.g., `os`). + 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, +} + +/// Find the "implicit" imports within the namespace package at the given path. +pub(crate) fn find(dir_path: &Path, exclusions: &[&Path]) -> HashMap { + let mut implicit_imports = HashMap::new(); + + // Enumerate all files and directories in the path, expanding links. + let Ok(entries) = fs::read_dir(dir_path) else { + return implicit_imports; + }; + + for entry in entries.flatten() { + let path = entry.path(); + + if exclusions.contains(&path.as_path()) { + continue; + } + + let Ok(file_type) = entry.file_type() else { + 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) + && !path + .with_extension(format!("{}.py", extension.to_str().unwrap())) + .exists() + && !path + .with_extension(format!("{}.pyi", extension.to_str().unwrap())) + .exists() + { + // E.g., `foo.abi3.so` becomes `foo`. + let file_stem = path + .file_stem() + .and_then(OsStr::to_str) + .and_then(|file_stem| { + file_stem.split_once('.').map(|(file_stem, _)| file_stem) + }); + let is_native_lib = true; + (file_stem, is_native_lib) + } else { + continue; + }; + + let Some(name) = file_stem else { + 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. + if implicit_imports + .get(&implicit_import.name) + .map_or(true, |implicit_import| !implicit_import.is_stub_file) + { + implicit_imports.insert(implicit_import.name.clone(), implicit_import); + } + } 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; + }; + + let implicit_import = ImplicitImport { + is_stub_file, + is_native_lib: false, + name: name.to_string(), + path: path.clone(), + py_typed: py_typed::get_py_typed_info(&path), + }; + implicit_imports.insert(implicit_import.name.clone(), implicit_import); + } + } + + implicit_imports +} + +/// Filter a map of implicit imports to only include those that were actually imported. +pub(crate) fn filter( + implicit_imports: &HashMap, + imported_symbols: &[String], +) -> Option> { + if implicit_imports.is_empty() || imported_symbols.is_empty() { + return None; + } + + let mut filtered_imports = HashMap::new(); + for implicit_import in implicit_imports.values() { + if imported_symbols.contains(&implicit_import.name) { + filtered_imports.insert(implicit_import.name.clone(), implicit_import.clone()); + } + } + + if filtered_imports.len() == implicit_imports.len() { + return None; + } + + Some(filtered_imports) +} diff --git a/crates/ruff_python_resolver/src/import_result.rs b/crates/ruff_python_resolver/src/import_result.rs new file mode 100644 index 0000000000..6ca7bd245c --- /dev/null +++ b/crates/ruff_python_resolver/src/import_result.rs @@ -0,0 +1,122 @@ +//! Interface that describes the output of the import resolver. + +use crate::implicit_imports::ImplicitImport; +use std::collections::HashMap; +use std::path::PathBuf; + +use crate::py_typed::PyTypedInfo; + +#[derive(Debug, Clone, PartialEq, Eq)] +#[allow(clippy::struct_excessive_bools)] +pub(crate) struct ImportResult { + /// Whether the import name was relative (e.g., ".foo"). + pub(crate) is_relative: bool, + + /// Whether the import was resolved to a file or module. + pub(crate) is_import_found: bool, + + /// The path was partially resolved, but the specific submodule + /// defining the import was not found. For example, `foo.bar` was + /// not found, but `foo` was found. + pub(crate) is_partly_resolved: bool, + + /// The import refers to a namespace package (i.e., a folder without + /// an `__init__.py[i]` file at the final level of resolution). By + /// convention, we insert empty `PathBuf` segments into the resolved + /// paths vector to indicate intermediary namespace packages. + pub(crate) is_namespace_package: bool, + + /// The final resolved directory contains an `__init__.py[i]` file. + pub(crate) is_init_file_present: bool, + + /// The import resolved to a stub (`.pyi`) file within a stub package. + pub(crate) is_stub_package: bool, + + /// The import resolved to a built-in, local, or third-party module. + pub(crate) import_type: ImportType, + + /// A vector of resolved absolute paths for each file in the module + /// name. Typically includes a sequence of `__init__.py` files, followed + /// by the Python file defining the import itself, though the exact + /// structure can vary. For example, namespace packages will be represented + /// by empty `PathBuf` segments in the vector. + /// + /// For example, resolving `import foo.bar` might yield `./foo/__init__.py` and `./foo/bar.py`, + /// or `./foo/__init__.py` and `./foo/bar/__init__.py`. + pub(crate) resolved_paths: Vec, + + /// The search path used to resolve the module. + pub(crate) search_path: Option, + + /// The resolved file is a type hint (i.e., a `.pyi` file), rather + /// than a Python (`.py`) file. + pub(crate) is_stub_file: bool, + + /// The resolved file is a native library. + pub(crate) is_native_lib: bool, + + /// The resolved file is a hint hint (i.e., a `.pyi` file) from + /// `typeshed` in the standard library. + pub(crate) is_stdlib_typeshed_file: bool, + + /// The resolved file is a hint hint (i.e., a `.pyi` file) from + /// `typeshed` in third-party stubs. + pub(crate) is_third_party_typeshed_file: bool, + + /// The resolved file is a type hint (i.e., a `.pyi` file) from + /// the configured typing directory. + pub(crate) is_local_typings_file: bool, + + /// A map from file to resolved path, for all implicitly imported + /// modules that are part of a namespace package. + pub(crate) implicit_imports: HashMap, + + /// Any implicit imports whose symbols were explicitly imported (i.e., via + /// a `from x import y` statement). + pub(crate) filtered_implicit_imports: HashMap, + + /// If the import resolved to a type hint (i.e., a `.pyi` file), then + /// a non-type-hint resolution will be stored here. + pub(crate) non_stub_import_result: Option>, + + /// Information extracted from the `py.typed` in the package used to + /// resolve the import, if any. + pub(crate) py_typed_info: Option, + + /// The directory of the package, if any. + pub(crate) package_directory: Option, +} + +impl ImportResult { + /// An import result that indicates that the import was not found. + pub(crate) fn not_found() -> Self { + Self { + is_relative: false, + is_import_found: false, + is_partly_resolved: false, + is_namespace_package: false, + is_init_file_present: false, + is_stub_package: false, + import_type: ImportType::Local, + resolved_paths: vec![], + search_path: None, + is_stub_file: false, + is_native_lib: false, + is_stdlib_typeshed_file: false, + is_third_party_typeshed_file: false, + is_local_typings_file: false, + implicit_imports: HashMap::default(), + filtered_implicit_imports: HashMap::default(), + non_stub_import_result: None, + py_typed_info: None, + package_directory: None, + } + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum ImportType { + BuiltIn, + ThirdParty, + Local, +} diff --git a/crates/ruff_python_resolver/src/lib.rs b/crates/ruff_python_resolver/src/lib.rs new file mode 100644 index 0000000000..c1fef3be66 --- /dev/null +++ b/crates/ruff_python_resolver/src/lib.rs @@ -0,0 +1,16 @@ +#![allow(dead_code)] + +mod config; +mod execution_environment; +mod host; +mod implicit_imports; +mod import_result; +mod module_descriptor; +mod native_module; +mod py_typed; +mod python_platform; +mod python_version; +mod resolver; +mod search; + +pub(crate) const SITE_PACKAGES: &str = "site-packages"; diff --git a/crates/ruff_python_resolver/src/module_descriptor.rs b/crates/ruff_python_resolver/src/module_descriptor.rs new file mode 100644 index 0000000000..7d71efafbc --- /dev/null +++ b/crates/ruff_python_resolver/src/module_descriptor.rs @@ -0,0 +1,16 @@ +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct ImportModuleDescriptor { + pub(crate) leading_dots: usize, + pub(crate) name_parts: Vec, + pub(crate) imported_symbols: Vec, +} + +impl ImportModuleDescriptor { + pub(crate) fn name(&self) -> String { + format!( + "{}{}", + ".".repeat(self.leading_dots), + &self.name_parts.join(".") + ) + } +} diff --git a/crates/ruff_python_resolver/src/native_module.rs b/crates/ruff_python_resolver/src/native_module.rs new file mode 100644 index 0000000000..110bde8500 --- /dev/null +++ b/crates/ruff_python_resolver/src/native_module.rs @@ -0,0 +1,14 @@ +//! Support for native Python extension modules. + +use std::ffi::OsStr; +use std::path::Path; + +/// Returns `true` if the given file extension is that of a native module. +pub(crate) fn is_native_module_file_extension(file_extension: &OsStr) -> bool { + file_extension == "so" || file_extension == "pyd" || file_extension == "dylib" +} + +/// Returns `true` if the given file name is that of a native module. +pub(crate) fn is_native_module_file_name(_module_name: &Path, _file_name: &Path) -> bool { + todo!() +} diff --git a/crates/ruff_python_resolver/src/py_typed.rs b/crates/ruff_python_resolver/src/py_typed.rs new file mode 100644 index 0000000000..258f801eed --- /dev/null +++ b/crates/ruff_python_resolver/src/py_typed.rs @@ -0,0 +1,40 @@ +//! Support for [PEP 561] (`py.typed` files). +//! +//! [PEP 561]: https://peps.python.org/pep-0561/ + +use std::path::{Path, PathBuf}; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct PyTypedInfo { + /// The path to the `py.typed` file. + py_typed_path: PathBuf, + + /// Whether the package is partially typed (as opposed to fully typed). + is_partially_typed: bool, +} + +/// Returns the `py.typed` information for the given directory, if any. +pub(crate) fn get_py_typed_info(dir_path: &Path) -> Option { + let py_typed_path = dir_path.join("py.typed"); + if py_typed_path.is_file() { + // Do a quick sanity check on the size before we attempt to read it. This + // file should always be really small - typically zero bytes in length. + let file_len = py_typed_path.metadata().ok()?.len(); + if file_len < 64 * 1024 { + // PEP 561 doesn't specify the format of "py.typed" in any detail other than + // to say that "If a stub package is partial it MUST include partial\n in a top + // level py.typed file." + let contents = std::fs::read_to_string(&py_typed_path).ok()?; + let is_partially_typed = + contents.contains("partial\n") || contents.contains("partial\r\n"); + Some(PyTypedInfo { + py_typed_path, + is_partially_typed, + }) + } else { + None + } + } else { + None + } +} diff --git a/crates/ruff_python_resolver/src/python_platform.rs b/crates/ruff_python_resolver/src/python_platform.rs new file mode 100644 index 0000000000..8ee2600518 --- /dev/null +++ b/crates/ruff_python_resolver/src/python_platform.rs @@ -0,0 +1,7 @@ +/// Enum to represent a Python platform. +#[derive(Debug, Copy, Clone, PartialEq, Eq)] +pub(crate) enum PythonPlatform { + Darwin, + Linux, + Windows, +} diff --git a/crates/ruff_python_resolver/src/python_version.rs b/crates/ruff_python_resolver/src/python_version.rs new file mode 100644 index 0000000000..aeb2a76b75 --- /dev/null +++ b/crates/ruff_python_resolver/src/python_version.rs @@ -0,0 +1,24 @@ +/// Enum to represent a Python version. +#[derive(Debug, Copy, Clone)] +pub(crate) enum PythonVersion { + Py37, + Py38, + Py39, + Py310, + Py311, + Py312, +} + +impl PythonVersion { + /// The directory name (e.g., in a virtual environment) for this Python version. + pub(crate) fn dir(self) -> &'static str { + match self { + PythonVersion::Py37 => "python3.7", + PythonVersion::Py38 => "python3.8", + PythonVersion::Py39 => "python3.9", + PythonVersion::Py310 => "python3.10", + PythonVersion::Py311 => "python3.11", + PythonVersion::Py312 => "python3.12", + } + } +} diff --git a/crates/ruff_python_resolver/src/resolver.rs b/crates/ruff_python_resolver/src/resolver.rs new file mode 100644 index 0000000000..08ff0bc09f --- /dev/null +++ b/crates/ruff_python_resolver/src/resolver.rs @@ -0,0 +1,1497 @@ +//! Resolves Python imports to their corresponding files on disk. + +use std::collections::HashMap; +use std::path::{Path, PathBuf}; + +use log::debug; + +use crate::config::Config; +use crate::execution_environment::ExecutionEnvironment; +use crate::implicit_imports::ImplicitImport; +use crate::import_result::{ImportResult, ImportType}; +use crate::module_descriptor::ImportModuleDescriptor; +use crate::search::get_typeshed_root; +use crate::{host, implicit_imports, native_module, py_typed, search}; + +#[allow(clippy::fn_params_excessive_bools)] +fn _resolve_absolute_import( + root: &Path, + module_descriptor: &ImportModuleDescriptor, + allow_partial: bool, + allow_native_lib: bool, + use_stub_package: bool, + allow_pyi: bool, + look_for_py_typed: bool, +) -> ImportResult { + if use_stub_package { + debug!("Attempting to resolve stub package using root path: {root:?}"); + } else { + debug!("Attempting to resolve using root path: {root:?}"); + } + + // Starting at the specified path, walk the file system to find the specified module. + let mut resolved_paths: Vec = Vec::new(); + let mut dir_path = root.to_path_buf(); + let mut is_namespace_package = false; + let mut is_init_file_present = false; + let mut is_stub_package = false; + let mut is_stub_file = false; + let mut is_native_lib = false; + let mut implicit_imports = HashMap::new(); + let mut package_directory = None; + let mut py_typed_info = None; + + // Ex) `from . import foo` + if module_descriptor.name_parts.is_empty() { + let py_file_path = dir_path.join("__init__.py"); + let pyi_file_path = dir_path.join("__init__.pyi"); + + if allow_pyi && pyi_file_path.is_file() { + debug!("Resolved import with file: {pyi_file_path:?}"); + resolved_paths.push(pyi_file_path.clone()); + } else if py_file_path.is_file() { + debug!("Resolved import with file: {py_file_path:?}"); + resolved_paths.push(py_file_path.clone()); + } else { + debug!("Partially resolved import with directory: {dir_path:?}"); + + // Add an empty path to indicate that the import is partially resolved. + resolved_paths.push(PathBuf::new()); + is_namespace_package = true; + } + + implicit_imports = implicit_imports::find(&dir_path, &[&py_file_path, &pyi_file_path]); + } else { + for (i, part) in module_descriptor.name_parts.iter().enumerate() { + let is_first_part = i == 0; + let is_last_part = i == module_descriptor.name_parts.len() - 1; + + // Extend the directory path with the next segment. + if use_stub_package && is_first_part { + dir_path = dir_path.join(format!("{part}-stubs")); + is_stub_package = true; + } else { + dir_path = dir_path.join(part); + } + + let found_directory = dir_path.is_dir(); + if found_directory { + if is_first_part { + package_directory = Some(dir_path.clone()); + } + + // Look for an `__init__.py[i]` in the directory. + let py_file_path = dir_path.join("__init__.py"); + let pyi_file_path = dir_path.join("__init__.pyi"); + is_init_file_present = false; + + if allow_pyi && pyi_file_path.is_file() { + debug!("Resolved import with file: {pyi_file_path:?}"); + resolved_paths.push(pyi_file_path.clone()); + if is_last_part { + is_stub_file = true; + } + is_init_file_present = true; + } else if py_file_path.is_file() { + debug!("Resolved import with file: {py_file_path:?}"); + resolved_paths.push(py_file_path.clone()); + is_init_file_present = true; + } + + if look_for_py_typed { + py_typed_info = + py_typed_info.or_else(|| py_typed::get_py_typed_info(&dir_path)); + } + + // We haven't reached the end of the import, and we found a matching directory. + // Proceed to the next segment. + if !is_last_part { + if !is_init_file_present { + resolved_paths.push(PathBuf::new()); + is_namespace_package = true; + py_typed_info = None; + } + continue; + } + + if is_init_file_present { + implicit_imports = + implicit_imports::find(&dir_path, &[&py_file_path, &pyi_file_path]); + break; + } + } + + // We couldn't find a matching directory, or the directory didn't contain an + // `__init__.py[i]` file. Look for an `.py[i]` file with the same name as the + // segment, in lieu of a directory. + let py_file_path = dir_path.with_extension("py"); + let pyi_file_path = dir_path.with_extension("pyi"); + + if allow_pyi && pyi_file_path.is_file() { + debug!("Resolved import with file: {pyi_file_path:?}"); + resolved_paths.push(pyi_file_path); + if is_last_part { + is_stub_file = true; + } + } else if py_file_path.is_file() { + debug!("Resolved import with file: {py_file_path:?}"); + resolved_paths.push(py_file_path); + } else { + if allow_native_lib && dir_path.is_dir() { + // We couldn't find a `.py[i]` file; search for a native library. + if let Some(native_lib_path) = dir_path + .read_dir() + .unwrap() + .flatten() + .filter(|entry| entry.file_type().map_or(false, |ft| ft.is_file())) + .find(|entry| { + native_module::is_native_module_file_name(&dir_path, &entry.path()) + }) + { + debug!("Resolved import with file: {native_lib_path:?}"); + is_native_lib = true; + resolved_paths.push(native_lib_path.path()); + } + } + + if !is_native_lib && found_directory { + debug!("Partially resolved import with directory: {dir_path:?}"); + resolved_paths.push(PathBuf::new()); + if is_last_part { + implicit_imports = + implicit_imports::find(&dir_path, &[&py_file_path, &pyi_file_path]); + is_namespace_package = true; + } + } else if is_native_lib { + debug!("Did not find file {py_file_path:?} or {pyi_file_path:?}"); + } + } + break; + } + } + + let import_found = if allow_partial { + !resolved_paths.is_empty() + } else { + resolved_paths.len() == module_descriptor.name_parts.len() + }; + + let is_partly_resolved = + !resolved_paths.is_empty() && resolved_paths.len() < module_descriptor.name_parts.len(); + + ImportResult { + is_relative: false, + is_import_found: import_found, + is_partly_resolved, + is_namespace_package, + is_init_file_present, + is_stub_package, + import_type: ImportType::Local, + resolved_paths, + search_path: Some(root.into()), + is_stub_file, + is_native_lib, + is_stdlib_typeshed_file: false, + is_third_party_typeshed_file: false, + is_local_typings_file: false, + implicit_imports, + filtered_implicit_imports: HashMap::default(), + non_stub_import_result: None, + py_typed_info, + package_directory, + } +} + +/// Resolve an absolute module import based on the import resolution algorithm +/// defined in [PEP 420]. +/// +/// [PEP 420]: https://peps.python.org/pep-0420/ +#[allow(clippy::fn_params_excessive_bools)] +fn resolve_absolute_import( + root: &Path, + module_descriptor: &ImportModuleDescriptor, + allow_partial: bool, + allow_native_lib: bool, + use_stub_package: bool, + allow_pyi: bool, + look_for_py_typed: bool, +) -> ImportResult { + if allow_pyi && use_stub_package { + // Search for packaged stubs first. PEP 561 indicates that package authors can ship + // stubs separately from the package implementation by appending `-stubs` to its + // top-level directory name. + let import_result = _resolve_absolute_import( + root, + module_descriptor, + allow_partial, + false, + true, + true, + true, + ); + + if import_result.package_directory.is_some() { + // If this is a namespace package that wasn't resolved, assume that + // it's a partial stub package and continue looking for a real package. + if !import_result.is_namespace_package || import_result.is_import_found { + return import_result; + } + } + } + + // Search for a "real" package. + _resolve_absolute_import( + root, + module_descriptor, + allow_partial, + allow_native_lib, + false, + allow_pyi, + look_for_py_typed, + ) +} + +/// Resolve an absolute module import based on the import resolution algorithm, +/// taking into account the various competing files to which the import could +/// resolve. +/// +/// For example, prefers local imports over third-party imports, and stubs over +/// non-stubs. +fn _resolve_best_absolute_import( + execution_environment: &ExecutionEnvironment, + module_descriptor: &ImportModuleDescriptor, + allow_pyi: bool, + config: &Config, + host: &Host, +) -> Option { + let import_name = module_descriptor.name(); + + // Search for local stub files (using `stub_path`). + if allow_pyi { + if let Some(stub_path) = config.stub_path.as_ref() { + debug!("Looking in stub path: {}", stub_path.display()); + + let mut typings_import = resolve_absolute_import( + stub_path, + module_descriptor, + false, + false, + true, + allow_pyi, + false, + ); + + if typings_import.is_import_found { + // Treat stub files as "local". + typings_import.import_type = ImportType::Local; + typings_import.is_local_typings_file = true; + + // If we resolved to a namespace package, ensure that all imported symbols are + // present in the namespace package's "implicit" imports. + if typings_import.is_namespace_package + && typings_import.resolved_paths[typings_import.resolved_paths.len() - 1] + .as_os_str() + .is_empty() + { + if _is_namespace_package_resolved( + module_descriptor, + &typings_import.implicit_imports, + ) { + return Some(typings_import); + } + } else { + return Some(typings_import); + } + } + + return None; + } + } + + // Look in the root directory of the execution environment. + debug!( + "Looking in root directory of execution environment: {}", + execution_environment.root.display() + ); + + let mut local_import = resolve_absolute_import( + &execution_environment.root, + module_descriptor, + false, + true, + true, + allow_pyi, + false, + ); + local_import.import_type = ImportType::Local; + + let mut best_result_so_far = Some(local_import); + + // Look in any extra paths. + for extra_path in &execution_environment.extra_paths { + debug!("Looking in extra path: {}", extra_path.display()); + + let mut local_import = resolve_absolute_import( + extra_path, + module_descriptor, + false, + true, + true, + allow_pyi, + false, + ); + local_import.import_type = ImportType::Local; + + best_result_so_far = Some(_pick_best_import( + best_result_so_far, + local_import, + module_descriptor, + )); + } + + // Look for third-party imports in Python's `sys` path. + for search_path in search::find_python_search_paths(config, host) { + debug!("Looking in Python search path: {}", search_path.display()); + + let mut third_party_import = resolve_absolute_import( + &search_path, + module_descriptor, + false, + true, + true, + allow_pyi, + true, + ); + third_party_import.import_type = ImportType::ThirdParty; + + best_result_so_far = Some(_pick_best_import( + best_result_so_far, + third_party_import, + module_descriptor, + )); + } + + // If a library is fully `py.typed`, prefer the current result. There's one exception: + // we're executing from `typeshed` itself. In that case, use the `typeshed` lookup below, + // rather than favoring `py.typed` libraries. + if let Some(typeshed_root) = get_typeshed_root(config, host) { + debug!( + "Looking in typeshed root directory: {}", + typeshed_root.display() + ); + if typeshed_root != execution_environment.root { + if best_result_so_far.as_ref().map_or(false, |result| { + result.py_typed_info.is_some() && !result.is_partly_resolved + }) { + return best_result_so_far; + } + } + } + + if allow_pyi && !module_descriptor.name_parts.is_empty() { + // Check for a stdlib typeshed file. + debug!("Looking for typeshed stdlib path: {}", import_name); + if let Some(mut typeshed_stdilib_import) = + _find_typeshed_path(module_descriptor, true, config, host) + { + typeshed_stdilib_import.is_stdlib_typeshed_file = true; + return Some(typeshed_stdilib_import); + } + + // Check for a third-party typeshed file. + debug!("Looking for typeshed third-party path: {}", import_name); + if let Some(mut typeshed_third_party_import) = + _find_typeshed_path(module_descriptor, false, config, host) + { + typeshed_third_party_import.is_third_party_typeshed_file = true; + + best_result_so_far = Some(_pick_best_import( + best_result_so_far, + typeshed_third_party_import, + module_descriptor, + )); + } + } + + // We weren't able to find an exact match, so return the best + // partial match. + 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: &HashMap, +) -> bool { + if !module_descriptor.imported_symbols.is_empty() { + // Pyright uses `!Array.from(moduleDescriptor.importedSymbols.keys()).some((symbol) => implicitImports.has(symbol))`. + // But 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. +/// +/// Supports both standard library and third-party `typeshed` lookups. +fn _find_typeshed_path( + module_descriptor: &ImportModuleDescriptor, + is_std_lib: bool, + config: &Config, + host: &Host, +) -> Option { + if is_std_lib { + debug!("Looking for typeshed `stdlib` path"); + } else { + debug!("Looking for typeshed `stubs` path"); + } + + let mut typeshed_paths = vec![]; + + if is_std_lib { + if let Some(path) = search::get_stdlib_typeshed_path(config, host) { + typeshed_paths.push(path); + } + } else { + if let Some(paths) = + search::get_third_party_typeshed_package_paths(module_descriptor, config, host) + { + typeshed_paths.extend(paths); + } + } + + for typeshed_path in typeshed_paths { + if typeshed_path.is_dir() { + let mut import_info = resolve_absolute_import( + &typeshed_path, + module_descriptor, + false, + false, + false, + true, + false, + ); + if import_info.is_import_found { + import_info.import_type = if is_std_lib { + ImportType::BuiltIn + } else { + ImportType::ThirdParty + }; + return Some(import_info); + } + } + } + + debug!("Typeshed path not found"); + None +} + +/// Given a current "best" import and a newly discovered result, returns the +/// preferred result. +fn _pick_best_import( + best_import_so_far: Option, + new_import: ImportResult, + module_descriptor: &ImportModuleDescriptor, +) -> ImportResult { + let Some(best_import_so_far) = best_import_so_far else { + return new_import; + }; + + if new_import.is_import_found { + // Prefer traditional over namespace packages. + let so_far_index = best_import_so_far + .resolved_paths + .iter() + .position(|path| !path.as_os_str().is_empty()); + let new_index = new_import + .resolved_paths + .iter() + .position(|path| !path.as_os_str().is_empty()); + if so_far_index != new_index { + match (so_far_index, new_index) { + (None, Some(_)) => return new_import, + (Some(_), None) => return best_import_so_far, + (Some(so_far_index), Some(new_index)) => { + return if so_far_index < new_index { + best_import_so_far + } else { + new_import + } + } + _ => {} + } + } + + // Prefer "found" over "not found". + if !best_import_so_far.is_import_found { + return new_import; + } + + // If both results are namespace imports, prefer the result that resolves all + // imported symbols. + if best_import_so_far.is_namespace_package && new_import.is_namespace_package { + if !module_descriptor.imported_symbols.is_empty() { + if !_is_namespace_package_resolved( + module_descriptor, + &best_import_so_far.implicit_imports, + ) { + if _is_namespace_package_resolved( + module_descriptor, + &new_import.implicit_imports, + ) { + return new_import; + } + + // Prefer the namespace package that has an `__init__.py[i]` file present in the + // final directory over one that does not. + if best_import_so_far.is_init_file_present && !new_import.is_init_file_present { + return best_import_so_far; + } + if !best_import_so_far.is_init_file_present && new_import.is_init_file_present { + return new_import; + } + } + } + } + + // Prefer "py.typed" over "non-py.typed". + if best_import_so_far.py_typed_info.is_some() && new_import.py_typed_info.is_none() { + return best_import_so_far; + } + if best_import_so_far.py_typed_info.is_none() && best_import_so_far.py_typed_info.is_some() + { + return new_import; + } + + // Prefer stub files (`.pyi`) over non-stub files (`.py`). + if best_import_so_far.is_stub_file && !new_import.is_stub_file { + return best_import_so_far; + } + if !best_import_so_far.is_stub_file && new_import.is_stub_file { + return new_import; + } + + // If we're still tied, prefer a shorter resolution path. + if best_import_so_far.resolved_paths.len() > new_import.resolved_paths.len() { + return new_import; + } + } else if new_import.is_partly_resolved { + let so_far_index = best_import_so_far + .resolved_paths + .iter() + .position(|path| !path.as_os_str().is_empty()); + let new_index = new_import + .resolved_paths + .iter() + .position(|path| !path.as_os_str().is_empty()); + if so_far_index != new_index { + match (so_far_index, new_index) { + (None, Some(_)) => return new_import, + (Some(_), None) => return best_import_so_far, + (Some(so_far_index), Some(new_index)) => { + return if so_far_index < new_index { + best_import_so_far + } else { + new_import + } + } + _ => {} + } + } + } + + best_import_so_far +} + +/// Resolve a relative import. +fn _resolve_relative_import( + source_file: &Path, + module_descriptor: &ImportModuleDescriptor, +) -> Option { + // Determine which search path this file is part of. + let mut directory = source_file; + for _ in 0..module_descriptor.leading_dots { + directory = directory.parent()?; + } + + // Now try to match the module parts from the current directory location. + let mut abs_import = resolve_absolute_import( + directory, + module_descriptor, + false, + true, + false, + true, + false, + ); + + if abs_import.is_stub_file { + // If we found a stub for a relative import, only search + // the same folder for the real module. Otherwise, it will + // error out on runtime. + abs_import.non_stub_import_result = Some(Box::new(resolve_absolute_import( + directory, + module_descriptor, + false, + true, + false, + false, + false, + ))); + } + + Some(abs_import) +} + +/// Resolve an absolute or relative import. +fn _resolve_import_strict( + source_file: &Path, + execution_environment: &ExecutionEnvironment, + module_descriptor: &ImportModuleDescriptor, + config: &Config, + host: &Host, +) -> ImportResult { + let import_name = module_descriptor.name(); + + if module_descriptor.leading_dots > 0 { + debug!("Resolving relative import for: {import_name}"); + + let relative_import = _resolve_relative_import(source_file, module_descriptor); + + if let Some(mut relative_import) = relative_import { + relative_import.is_relative = true; + return relative_import; + } + } else { + debug!("Resolving best absolute import for: {import_name}"); + + let best_import = _resolve_best_absolute_import( + execution_environment, + module_descriptor, + true, + config, + host, + ); + + if let Some(mut best_import) = best_import { + if best_import.is_stub_file { + debug!("Resolving best non-stub absolute import for: {import_name}"); + + best_import.non_stub_import_result = Some(Box::new( + _resolve_best_absolute_import( + execution_environment, + module_descriptor, + false, + config, + host, + ) + .unwrap_or_else(ImportResult::not_found), + )); + } + return best_import; + } + } + + ImportResult::not_found() +} + +/// Resolves an import, given the current file and the import descriptor. +fn resolve_import( + source_file: &Path, + execution_environment: &ExecutionEnvironment, + module_descriptor: &ImportModuleDescriptor, + config: &Config, + host: &Host, +) -> ImportResult { + let import_result = _resolve_import_strict( + source_file, + execution_environment, + module_descriptor, + config, + host, + ); + if import_result.is_import_found || module_descriptor.leading_dots > 0 { + return import_result; + } + + // If we weren't able to resolve an absolute import, try resolving it in the + // importing file's directory, then the parent directory, and so on, until the + // import root is reached. + let root = execution_environment.root.as_path(); + if source_file.starts_with(root) { + let mut current = source_file; + while let Some(parent) = current.parent() { + if parent == root { + break; + } + + debug!("Resolving absolute import in parent: {}", parent.display()); + + let mut result = resolve_absolute_import( + parent, + module_descriptor, + false, + false, + false, + true, + false, + ); + + if result.is_import_found { + if let Some(implicit_imports) = implicit_imports::filter( + &result.implicit_imports, + &module_descriptor.imported_symbols, + ) { + result.implicit_imports = implicit_imports; + } + return result; + } + + current = parent; + } + } + + ImportResult::not_found() +} + +#[cfg(test)] +mod tests { + use std::fs::{create_dir_all, File}; + use std::io::{self, Write}; + use std::path::{Path, PathBuf}; + + use log::debug; + use tempfile::TempDir; + + use crate::config::Config; + use crate::execution_environment::ExecutionEnvironment; + use crate::import_result::{ImportResult, ImportType}; + use crate::module_descriptor::ImportModuleDescriptor; + use crate::python_platform::PythonPlatform; + use crate::python_version::PythonVersion; + use crate::resolver::resolve_import; + use crate::{host, SITE_PACKAGES}; + + struct PythonFile { + path: PathBuf, + content: String, + } + + impl PythonFile { + fn new(path: impl Into, content: impl Into) -> Self { + Self { + path: path.into(), + content: content.into(), + } + } + + fn create(&self, dir: impl AsRef) -> io::Result { + let file_path = dir.as_ref().join(self.path.as_path()); + if let Some(parent) = file_path.parent() { + create_dir_all(parent)?; + } + let mut f = File::create(&file_path)?; + f.write_all(self.content.as_bytes())?; + f.sync_all()?; + + Ok(file_path) + } + } + + fn _resolve>( + source_file: impl AsRef, + name: &str, + root: T, + extra_paths: Vec, + library: Option, + stub_path: Option, + typeshed_path: Option, + ) -> ImportResult { + let execution_environment = ExecutionEnvironment { + root: root.into(), + python_version: PythonVersion::Py37, + python_platform: PythonPlatform::Darwin, + extra_paths, + }; + + let module_descriptor = ImportModuleDescriptor { + leading_dots: name.chars().take_while(|c| *c == '.').count(), + name_parts: name + .chars() + .skip_while(|c| *c == '.') + .collect::() + .split('.') + .map(std::string::ToString::to_string) + .collect(), + imported_symbols: Vec::new(), + }; + + let config = Config { + venv_path: None, + venv: None, + python_path: None, + typeshed_path: typeshed_path.map(Into::into), + stub_path: stub_path.map(Into::into), + default_python_version: None, + }; + + let host = host::StaticHost::new(if let Some(library) = library { + vec![library.into()] + } else { + Vec::new() + }); + + resolve_import( + source_file.as_ref(), + &execution_environment, + &module_descriptor, + &config, + &host, + ) + } + + fn resolve>( + source_file: impl AsRef, + name: &str, + root: T, + ) -> ImportResult { + _resolve(source_file, name, root, Vec::new(), None, None, None) + } + + fn resolve_with_extra_paths>( + source_file: impl AsRef, + name: &str, + root: T, + extra_paths: Vec, + ) -> ImportResult { + _resolve(source_file, name, root, extra_paths, None, None, None) + } + + fn resolve_with_library>( + source_file: impl AsRef, + name: &str, + root: T, + library: T, + ) -> ImportResult { + _resolve( + source_file, + name, + root, + Vec::new(), + Some(library), + None, + None, + ) + } + + fn resolve_with_stub_path>( + source_file: impl AsRef, + name: &str, + root: T, + library: T, + stub_path: T, + ) -> ImportResult { + _resolve( + source_file, + name, + root, + Vec::new(), + Some(library), + Some(stub_path), + None, + ) + } + + fn resolve_with_typeshed_path>( + source_file: impl AsRef, + name: &str, + root: T, + library: T, + typeshed_path: T, + ) -> ImportResult { + _resolve( + source_file, + name, + root, + Vec::new(), + Some(library), + None, + Some(typeshed_path), + ) + } + + #[test] + fn partial_stub_file_exists() -> io::Result<()> { + env_logger::builder().is_test(true).try_init().ok(); + + let dir = TempDir::new()?; + let library = TempDir::new()?.path().join("lib").join(SITE_PACKAGES); + + let partial_stub_pyi = + PythonFile::new("myLib-stubs/partialStub.pyi", "def test(): ...").create(&library)?; + PythonFile::new("myLib-stubs/py.typed", "partial\n").create(&library)?; + let partial_stub_py = + PythonFile::new("myLib/partialStub.py", "def test(): ...").create(&library)?; + + let result = resolve_with_library( + partial_stub_py, + "myLib.partialStub", + dir.path(), + library.as_path(), + ); + + assert!(result.is_import_found); + assert!(result.is_stub_file); + assert_eq!(result.import_type, ImportType::ThirdParty); + assert_eq!( + result.resolved_paths, + // TODO(charlie): Pyright matches on `libraryRoot, 'myLib', 'partialStub.pyi'` here. + // But that file doesn't exist. There's some kind of transform. + vec![PathBuf::new(), partial_stub_pyi] + ); + + Ok(()) + } + + #[test] + fn partial_stub_init_exists() -> io::Result<()> { + env_logger::builder().is_test(true).try_init().ok(); + + let dir = TempDir::new()?; + let library = TempDir::new()?.path().join("lib").join(SITE_PACKAGES); + + PythonFile::new("myLib-stubs/py.typed", "partial\n").create(&library)?; + let partial_stub_init_pyi = + PythonFile::new("myLib-stubs/__init__.pyi", "def test(): ...").create(&library)?; + let partial_stub_init_py = + PythonFile::new("myLib/__init__.py", "def test(): ...").create(&library)?; + + let result = + resolve_with_library(partial_stub_init_py, "myLib", dir.path(), library.as_path()); + + assert!(result.is_import_found); + assert!(result.is_stub_file); + assert_eq!(result.import_type, ImportType::ThirdParty); + assert_eq!( + result.resolved_paths, + // TODO(charlie): Pyright matches on `libraryRoot, 'myLib', '__init__.pyi'` here. + // But that file doesn't exist. There's some kind of transform. + vec![partial_stub_init_pyi] + ); + + Ok(()) + } + + #[test] + fn side_by_side_files() -> io::Result<()> { + let dir = TempDir::new()?; + let library = TempDir::new()?.path().join("lib").join(SITE_PACKAGES); + + let my_file = PythonFile::new("myFile.py", "# not used").create(dir.path())?; + let side_by_side_stub_file = + PythonFile::new("myLib-stubs/partialStub.pyi", "def test(): ...").create(&library)?; + PythonFile::new("myLib-stubs/py.typed", "partial\n").create(&library)?; + PythonFile::new("myLib/partialStub.pyi", "# empty").create(&library)?; + PythonFile::new("myLib/partialStub.py", "def test(): pass").create(&library)?; + let partial_stub_file = + PythonFile::new("myLib-stubs/partialStub2.pyi", "def test(): ...").create(&library)?; + PythonFile::new("myLib/partialStub2.py", "def test(): pass").create(&library)?; + + // Stub package wins over original package (per PEP 561 rules). + let side_by_side_result = + resolve_with_library(&my_file, "myLib.partialStub", dir.path(), &library); + assert!(side_by_side_result.is_import_found); + assert!(side_by_side_result.is_stub_file); + assert_eq!( + side_by_side_result.resolved_paths, + vec![PathBuf::new(), side_by_side_stub_file] + ); + + // Side by side stub doesn't completely disable partial stub. + let partial_stub_result = + resolve_with_library(&my_file, "myLib.partialStub2", dir.path(), &library); + assert!(partial_stub_result.is_import_found); + assert!(partial_stub_result.is_stub_file); + assert_eq!( + partial_stub_result.resolved_paths, + vec![PathBuf::new(), partial_stub_file] + ); + + Ok(()) + } + + #[test] + fn stub_package() -> io::Result<()> { + env_logger::builder().is_test(true).try_init().ok(); + + let dir = TempDir::new()?; + let library = TempDir::new()?.path().join("lib").join(SITE_PACKAGES); + + PythonFile::new("myLib-stubs/stub.pyi", "# empty").create(&library)?; + PythonFile::new("myLib-stubs/__init__.pyi", "# empty").create(&library)?; + let partial_stub_py = + PythonFile::new("myLib/partialStub.py", "def test(): ...").create(&library)?; + + let result = resolve_with_library( + partial_stub_py, + "myLib.partialStub", + dir.path(), + library.as_path(), + ); + + // If fully typed stub package exists, that wins over the real package. + assert!(!result.is_import_found); + + Ok(()) + } + + #[test] + fn stub_namespace_package() -> io::Result<()> { + env_logger::builder().is_test(true).try_init().ok(); + + let dir = TempDir::new()?; + let library = TempDir::new()?.path().join("lib").join(SITE_PACKAGES); + + PythonFile::new("myLib-stubs/stub.pyi", "# empty").create(&library)?; + let partial_stub_py = + PythonFile::new("myLib/partialStub.py", "def test(): ...").create(&library)?; + + let result = resolve_with_library( + partial_stub_py.clone(), + "myLib.partialStub", + dir.path(), + library.as_path(), + ); + + // If fully typed stub package exists, that wins over the real package. + assert!(result.is_import_found); + assert!(!result.is_stub_file); + assert_eq!(result.resolved_paths, vec![PathBuf::new(), partial_stub_py]); + + Ok(()) + } + + #[test] + fn stub_in_typing_folder_over_partial_stub_package() -> io::Result<()> { + env_logger::builder().is_test(true).try_init().ok(); + + let dir = TempDir::new()?; + let typing_folder = dir.path().join("typing"); + let library = TempDir::new()?.path().join("lib").join(SITE_PACKAGES); + + PythonFile::new("myLib-stubs/py.typed", "partial\n").create(&library)?; + PythonFile::new("myLib-stubs/__init__.pyi", "").create(&library)?; + let my_lib_pyi = PythonFile::new("myLib.pyi", "# empty").create(&typing_folder)?; + let my_lib_init_py = + PythonFile::new("myLib/__init__.py", "def test(): pass").create(&library)?; + + let result = resolve_with_stub_path( + my_lib_init_py, + "myLib", + dir.path(), + library.as_path(), + typing_folder.as_path(), + ); + + // If the package exists in typing folder, that gets picked up first (so we resolve to + // `myLib.pyi`). + assert!(result.is_import_found); + assert!(result.is_stub_file); + assert_eq!(result.resolved_paths, vec![my_lib_pyi]); + + Ok(()) + } + + #[test] + fn partial_stub_package_in_typing_folder() -> io::Result<()> { + env_logger::builder().is_test(true).try_init().ok(); + + let dir = TempDir::new()?; + let typing_folder = dir.path().join("typing"); + let library = TempDir::new()?.path().join("lib").join(SITE_PACKAGES); + + PythonFile::new("myLib-stubs/py.typed", "partial\n").create(&typing_folder)?; + let my_lib_stubs_init_pyi = PythonFile::new("myLib-stubs/__init__.pyi", "def test(): ...") + .create(&typing_folder)?; + let my_lib_init_py = + PythonFile::new("myLib/__init__.py", "def test(): pass").create(&library)?; + + let result = resolve_with_stub_path( + my_lib_init_py, + "myLib", + dir.path(), + library.as_path(), + typing_folder.as_path(), + ); + + // If the package exists in typing folder, that gets picked up first (so we resolve to + // `myLib.pyi`). + assert!(result.is_import_found); + assert!(result.is_stub_file); + assert_eq!(result.resolved_paths, vec![my_lib_stubs_init_pyi]); + + Ok(()) + } + + #[test] + fn typeshed_folder() -> io::Result<()> { + env_logger::builder().is_test(true).try_init().ok(); + + let dir = TempDir::new()?; + let library = TempDir::new()?.path().join("lib").join(SITE_PACKAGES); + let typeshed_folder = dir.path().join("ts"); + + PythonFile::new("myLib-stubs/py.typed", "partial\n").create(&library)?; + let my_lib_stubs_init_pyi = + PythonFile::new("myLib-stubs/__init__.pyi", "def test(): ...").create(&library)?; + PythonFile::new("stubs/myLibPackage/myLib.pyi", "# empty").create(&typeshed_folder)?; + let my_lib_init_py = + PythonFile::new("myLib/__init__.py", "def test(): pass").create(&library)?; + + let result = resolve_with_typeshed_path( + my_lib_init_py, + "myLib", + dir.path(), + library.as_path(), + typeshed_folder.as_path(), + ); + + // Stub packages win over typeshed. + assert!(result.is_import_found); + assert!(result.is_stub_file); + assert_eq!(result.resolved_paths, vec![my_lib_stubs_init_pyi]); + + Ok(()) + } + + #[test] + fn py_typed_file() -> io::Result<()> { + env_logger::builder().is_test(true).try_init().ok(); + + let dir = TempDir::new()?; + let library = TempDir::new()?.path().join("lib").join(SITE_PACKAGES); + + let partial_stub_init_pyi = + PythonFile::new("myLib-stubs/__init__.pyi", "def test(): ...").create(&library)?; + PythonFile::new("myLib-stubs/py.typed", "partial\n").create(&library)?; + PythonFile::new("myLib/__init__.py", "def test(): pass").create(&library)?; + let package_py_typed = PythonFile::new("myLib/py.typed", "# typed").create(&library)?; + + let result = resolve_with_library(package_py_typed, "myLib", dir.path(), library.as_path()); + + // Partial stub package always overrides original package. + assert!(result.is_import_found); + assert!(result.is_stub_file); + assert_eq!(result.resolved_paths, vec![partial_stub_init_pyi]); + + Ok(()) + } + + #[test] + fn py_typed_library() -> io::Result<()> { + env_logger::builder().is_test(true).try_init().ok(); + + let dir = TempDir::new()?; + let library = TempDir::new()?.path().join("lib").join(SITE_PACKAGES); + let typeshed_folder = dir.path().join("ts"); + + let init_py = PythonFile::new("os/__init__.py", "def test(): ...").create(&library)?; + PythonFile::new("os/py.typed", "").create(&library)?; + let typeshed_init_pyi = + PythonFile::new("stubs/os/os/__init__.pyi", "# empty").create(&typeshed_folder)?; + + let result = resolve_with_typeshed_path( + typeshed_init_pyi, + "os", + dir.path(), + library.as_path(), + typeshed_folder.as_path(), + ); + + assert!(result.is_import_found); + assert_eq!(result.resolved_paths, vec![init_py]); + + Ok(()) + } + + #[test] + fn non_py_typed_library() -> io::Result<()> { + env_logger::builder().is_test(true).try_init().ok(); + + let dir = TempDir::new()?; + let library = TempDir::new()?.path().join("lib").join(SITE_PACKAGES); + let typeshed_folder = dir.path().join("ts"); + + PythonFile::new("os/__init__.py", "def test(): ...").create(&library)?; + let typeshed_init_pyi = + PythonFile::new("stubs/os/os/__init__.pyi", "# empty").create(&typeshed_folder)?; + + let result = resolve_with_typeshed_path( + typeshed_init_pyi.clone(), + "os", + dir.path(), + library.as_path(), + typeshed_folder.as_path(), + ); + + assert!(result.is_import_found); + assert_eq!(result.import_type, ImportType::ThirdParty); + assert_eq!(result.resolved_paths, vec![typeshed_init_pyi]); + + Ok(()) + } + + #[test] + fn import_side_by_side_file_root() -> io::Result<()> { + env_logger::builder().is_test(true).try_init().ok(); + + let dir = TempDir::new()?; + + let file1 = PythonFile::new("file1.py", "import file1").create(&dir)?; + let file2 = PythonFile::new("file2.py", "import file2").create(&dir)?; + + let result = resolve(file2, "file1", dir.path()); + + assert!(result.is_import_found); + assert_eq!(result.import_type, ImportType::Local); + assert_eq!(result.resolved_paths, vec![file1]); + + Ok(()) + } + + #[test] + fn import_side_by_side_file_sub_folder() -> io::Result<()> { + env_logger::builder().is_test(true).try_init().ok(); + + let dir = TempDir::new()?; + + let test_init = PythonFile::new("test/__init__.py", "").create(&dir)?; + let test_file1 = PythonFile::new("test/file1.py", "import file1").create(&dir)?; + let test_file2 = PythonFile::new("test/file2.py", "import file2").create(&dir)?; + + let result = resolve(test_file2, "test.file1", dir.path()); + + assert!(result.is_import_found); + assert_eq!(result.import_type, ImportType::Local); + assert_eq!(result.resolved_paths, vec![test_init, test_file1]); + + Ok(()) + } + + #[test] + fn import_side_by_side_file_sub_under_src_folder() -> io::Result<()> { + env_logger::builder().is_test(true).try_init().ok(); + + let dir = TempDir::new()?; + + let nested_init = PythonFile::new("src/nested/__init__.py", "").create(&dir)?; + let nested_file1 = PythonFile::new("src/nested/file1.py", "import file1").create(&dir)?; + let nested_file2 = PythonFile::new("src/nested/file2.py", "import file2").create(&dir)?; + + let result = resolve(nested_file2, "nested.file1", dir.path()); + + assert!(result.is_import_found); + assert_eq!(result.import_type, ImportType::Local); + assert_eq!(result.resolved_paths, vec![nested_init, nested_file1]); + + Ok(()) + } + + #[test] + fn import_file_sub_under_containing_folder() -> io::Result<()> { + env_logger::builder().is_test(true).try_init().ok(); + + let dir = TempDir::new()?; + + let nested_file1 = + PythonFile::new("src/nested/file1.py", "def test1(): ... ").create(&dir)?; + let nested_file2 = + PythonFile::new("src/nested/nested2/file2.py", "def test2(): ...").create(&dir)?; + + let result = resolve(nested_file2, "file1", dir.path()); + + assert!(result.is_import_found); + assert_eq!(result.import_type, ImportType::Local); + assert_eq!(result.resolved_paths, vec![nested_file1]); + + Ok(()) + } + + #[test] + fn import_side_by_side_file_sub_under_lib_folder() -> io::Result<()> { + env_logger::builder().is_test(true).try_init().ok(); + + let dir = TempDir::new()?; + let library = TempDir::new()?.path().join("lib").join(SITE_PACKAGES); + + PythonFile::new("myLib/file1.py", "def test1(): ...").create(&library)?; + let file2 = PythonFile::new("myLib/file2.py", "def test2(): ...").create(&library)?; + + let result = resolve(file2, "file1", dir.path()); + + debug!("result: {:?}", result); + + assert!(!result.is_import_found); + + Ok(()) + } + + #[test] + fn nested_namespace_package_1() -> io::Result<()> { + // See: https://github.com/microsoft/pyright/issues/5089. + env_logger::builder().is_test(true).try_init().ok(); + + let dir = TempDir::new()?; + + let file = PythonFile::new("package1/a/b/c/d.py", "def f(): pass").create(&dir)?; + let package1_init = PythonFile::new("package1/a/__init__.py", "").create(&dir)?; + let package2_init = PythonFile::new("package2/a/__init__.py", "").create(&dir)?; + + let package1 = dir.path().join("package1"); + let package2 = dir.path().join("package2"); + + let result = resolve_with_extra_paths( + package2_init, + "a.b.c.d", + dir.path(), + vec![package1, package2], + ); + + assert!(result.is_import_found); + assert_eq!(result.import_type, ImportType::Local); + assert_eq!( + result.resolved_paths, + vec![package1_init, PathBuf::new(), PathBuf::new(), file] + ); + + Ok(()) + } + + #[test] + fn nested_namespace_package_2() -> io::Result<()> { + // See: https://github.com/microsoft/pyright/issues/5089. + env_logger::builder().is_test(true).try_init().ok(); + + let dir = TempDir::new()?; + + let file = PythonFile::new("package1/a/b/c/d.py", "def f(): pass").create(&dir)?; + let package1_init = PythonFile::new("package1/a/b/c/__init__.py", "").create(&dir)?; + let package2_init = PythonFile::new("package2/a/b/c/__init__.py", "").create(&dir)?; + + let package1 = dir.path().join("package1"); + let package2 = dir.path().join("package2"); + + let result = resolve_with_extra_paths( + package2_init, + "a.b.c.d", + dir.path(), + vec![package1, package2], + ); + + assert!(result.is_import_found); + assert_eq!(result.import_type, ImportType::Local); + assert_eq!( + result.resolved_paths, + vec![PathBuf::new(), PathBuf::new(), package1_init, file] + ); + + Ok(()) + } + + #[test] + fn nested_namespace_package_3() -> io::Result<()> { + // See: https://github.com/microsoft/pyright/issues/5089. + env_logger::builder().is_test(true).try_init().ok(); + + let dir = TempDir::new()?; + + PythonFile::new("package1/a/b/c/d.py", "def f(): pass").create(&dir)?; + let package2_init = PythonFile::new("package2/a/__init__.py", "").create(&dir)?; + + let package1 = dir.path().join("package1"); + let package2 = dir.path().join("package2"); + + let result = resolve_with_extra_paths( + package2_init, + "a.b.c.d", + dir.path(), + vec![package1, package2], + ); + + assert!(!result.is_import_found); + + Ok(()) + } + + #[test] + fn nested_namespace_package_4() -> io::Result<()> { + // See: https://github.com/microsoft/pyright/issues/5089. + env_logger::builder().is_test(true).try_init().ok(); + + let dir = TempDir::new()?; + + PythonFile::new("package1/a/b/__init__.py", "").create(&dir)?; + PythonFile::new("package1/a/b/c.py", "def f(): pass").create(&dir)?; + PythonFile::new("package2/a/__init__.py", "").create(&dir)?; + let package2_a_b_init = PythonFile::new("package2/a/b/__init__.py", "").create(&dir)?; + + let package1 = dir.path().join("package1"); + let package2 = dir.path().join("package2"); + + let result = resolve_with_extra_paths( + package2_a_b_init, + "a.b.c", + dir.path(), + vec![package1, package2], + ); + + assert!(!result.is_import_found); + + Ok(()) + } + + // New tests, don't exist upstream. + #[test] + fn relative_import_side_by_side_file_root() -> io::Result<()> { + env_logger::builder().is_test(true).try_init().ok(); + + let dir = TempDir::new()?; + + let file1 = PythonFile::new("file1.py", "def test(): ...").create(&dir)?; + let file2 = PythonFile::new("file2.py", "def test(): ...").create(&dir)?; + + let result = resolve(file2, ".file1", dir.path()); + + assert!(result.is_import_found); + assert_eq!(result.import_type, ImportType::Local); + assert_eq!(result.resolved_paths, vec![file1]); + + Ok(()) + } + + #[test] + fn invalid_relative_import_side_by_side_file_root() -> io::Result<()> { + env_logger::builder().is_test(true).try_init().ok(); + + let dir = TempDir::new()?; + + PythonFile::new("file1.py", "def test(): ...").create(&dir)?; + let file2 = PythonFile::new("file2.py", "def test(): ...").create(&dir)?; + + let result = resolve(file2, "..file1", dir.path()); + + assert!(!result.is_import_found); + + Ok(()) + } +} diff --git a/crates/ruff_python_resolver/src/search.rs b/crates/ruff_python_resolver/src/search.rs new file mode 100644 index 0000000000..66145aef78 --- /dev/null +++ b/crates/ruff_python_resolver/src/search.rs @@ -0,0 +1,282 @@ +//! Determine the appropriate search paths for the Python environment. + +use std::collections::HashMap; +use std::ffi::OsStr; +use std::fs; +use std::path::{Path, PathBuf}; + +use log::debug; + +use crate::config::Config; +use crate::module_descriptor::ImportModuleDescriptor; +use crate::python_version::PythonVersion; +use crate::{host, SITE_PACKAGES}; + +/// Find the `site-packages` directory for the specified Python version. +fn find_site_packages_path( + lib_path: &Path, + python_version: Option, +) -> Option { + if lib_path.is_dir() { + debug!( + "Found path `{}`; looking for site-packages", + lib_path.display() + ); + } else { + debug!("Did not find `{}`", lib_path.display()); + } + + let site_packages_path = lib_path.join(SITE_PACKAGES); + if site_packages_path.is_dir() { + debug!("Found path `{}`", site_packages_path.display()); + return Some(site_packages_path); + } + + debug!( + "Did not find `{}`, so looking for Python subdirectory", + site_packages_path.display() + ); + + // There's no `site-packages` directory in the library directory; look for a `python3.X` + // directory instead. + let candidate_dirs: Vec = fs::read_dir(lib_path) + .ok()? + .filter_map(|entry| { + let entry = entry.ok()?; + let metadata = entry.metadata().ok()?; + + if metadata.file_type().is_dir() { + let dir_path = entry.path(); + if dir_path + .file_name() + .and_then(OsStr::to_str)? + .starts_with("python3.") + { + if dir_path.join(SITE_PACKAGES).is_dir() { + return Some(dir_path); + } + } + } else if metadata.file_type().is_symlink() { + let symlink_path = fs::read_link(entry.path()).ok()?; + if symlink_path + .file_name() + .and_then(OsStr::to_str)? + .starts_with("python3.") + { + if symlink_path.join(SITE_PACKAGES).is_dir() { + return Some(symlink_path); + } + } + } + + None + }) + .collect(); + + // If a `python3.X` directory does exist (and `3.X` matches the current Python version), + // prefer it over any other Python directories. + if let Some(python_version) = python_version { + if let Some(preferred_dir) = candidate_dirs.iter().find(|dir| { + dir.file_name() + .and_then(OsStr::to_str) + .map_or(false, |name| name == python_version.dir()) + }) { + debug!("Found path `{}`", preferred_dir.display()); + return Some(preferred_dir.join(SITE_PACKAGES)); + } + } + + // Fallback to the first `python3.X` directory that we found. + let default_dir = candidate_dirs.first()?; + debug!("Found path `{}`", default_dir.display()); + Some(default_dir.join(SITE_PACKAGES)) +} + +fn get_paths_from_pth_files(parent_dir: &Path) -> Vec { + fs::read_dir(parent_dir) + .unwrap() + .flatten() + .filter(|entry| { + // Collect all *.pth files. + let Ok(file_type) = entry.file_type() else { + return false; + }; + file_type.is_file() || file_type.is_symlink() + }) + .map(|entry| entry.path()) + .filter(|path| path.extension() == Some(OsStr::new("pth"))) + .filter(|path| { + // Skip all files that are much larger than expected. + let Ok(metadata) = path.metadata() else { + return false; + }; + let file_len = metadata.len(); + file_len > 0 && file_len < 64 * 1024 + }) + .filter_map(|path| { + let data = fs::read_to_string(&path).ok()?; + for line in data.lines() { + let trimmed_line = line.trim(); + if !trimmed_line.is_empty() + && !trimmed_line.starts_with('#') + && !trimmed_line.starts_with("import") + { + let pth_path = parent_dir.join(trimmed_line); + if pth_path.is_dir() { + return Some(pth_path); + } + } + } + None + }) + .collect() +} + +/// Find the Python search paths for the given virtual environment. +pub(crate) fn find_python_search_paths( + config: &Config, + host: &Host, +) -> Vec { + if let Some(venv_path) = config.venv_path.as_ref() { + if let Some(venv) = config.venv.as_ref() { + let mut found_paths = vec![]; + + for lib_name in ["lib", "Lib", "lib64"] { + let lib_path = venv_path.join(venv).join(lib_name); + if let Some(site_packages_path) = + find_site_packages_path(&lib_path, config.default_python_version) + { + // Add paths from any `.pth` files in each of the `site-packages` directories. + found_paths.extend(get_paths_from_pth_files(&site_packages_path)); + + // Add the `site-packages` directory to the search path. + found_paths.push(site_packages_path); + } + } + + if !found_paths.is_empty() { + found_paths.sort(); + found_paths.dedup(); + + debug!("Found the following `site-packages` dirs"); + for path in &found_paths { + debug!(" {}", path.display()); + } + + return found_paths; + } + } + } + + // Fall back to the Python interpreter. + host.python_search_paths() +} + +/// Determine the relevant Python search paths. +fn get_python_search_paths(config: &Config, host: &Host) -> Vec { + // TODO(charlie): Cache search paths. + find_python_search_paths(config, host) +} + +/// Determine the root of the `typeshed` directory. +pub(crate) fn get_typeshed_root(config: &Config, host: &Host) -> Option { + if let Some(typeshed_path) = config.typeshed_path.as_ref() { + // Did the user specify a typeshed path? + if typeshed_path.is_dir() { + return Some(typeshed_path.clone()); + } + } else { + // If not, we'll look in the Python search paths. + for python_search_path in get_python_search_paths(config, host) { + let possible_typeshed_path = python_search_path.join("typeshed"); + if possible_typeshed_path.is_dir() { + return Some(possible_typeshed_path); + } + } + } + + None +} + +/// Format the expected `typeshed` subdirectory. +fn format_typeshed_subdirectory(typeshed_path: &Path, is_stdlib: bool) -> PathBuf { + typeshed_path.join(if is_stdlib { "stdlib" } else { "stubs" }) +} + +/// Determine the current `typeshed` subdirectory. +fn get_typeshed_subdirectory( + is_stdlib: bool, + config: &Config, + host: &Host, +) -> Option { + let typeshed_path = get_typeshed_root(config, host)?; + let typeshed_path = format_typeshed_subdirectory(&typeshed_path, is_stdlib); + if typeshed_path.is_dir() { + Some(typeshed_path) + } else { + None + } +} + +/// Determine the current `typeshed` subdirectory for the standard library. +pub(crate) fn get_stdlib_typeshed_path( + config: &Config, + host: &Host, +) -> Option { + get_typeshed_subdirectory(true, config, host) +} + +/// Generate a map from PyPI-registered package name to a list of paths +/// containing the package's stubs. +fn build_typeshed_third_party_package_map(third_party_dir: &Path) -> HashMap> { + let mut package_map = HashMap::new(); + + // Iterate over every directory. + for outer_entry in fs::read_dir(third_party_dir).unwrap() { + let outer_entry = outer_entry.unwrap(); + if outer_entry.file_type().unwrap().is_dir() { + // Iterate over any subdirectory children. + for inner_entry in fs::read_dir(outer_entry.path()).unwrap() { + let inner_entry = inner_entry.unwrap(); + + if inner_entry.file_type().unwrap().is_dir() { + package_map + .entry(inner_entry.file_name().to_string_lossy().to_string()) + .or_insert_with(Vec::new) + .push(outer_entry.path()); + } else if inner_entry.file_type().unwrap().is_file() { + if inner_entry + .path() + .extension() + .map_or(false, |extension| extension == "pyi") + { + let stripped_file_name = inner_entry + .path() + .file_stem() + .unwrap() + .to_string_lossy() + .to_string(); + package_map + .entry(stripped_file_name) + .or_insert_with(Vec::new) + .push(outer_entry.path()); + } + } + } + } + } + + package_map +} + +/// Determine the current `typeshed` subdirectory for a third-party package. +pub(crate) fn get_third_party_typeshed_package_paths( + module_descriptor: &ImportModuleDescriptor, + config: &Config, + host: &Host, +) -> Option> { + let typeshed_path = get_typeshed_subdirectory(false, config, host)?; + let package_paths = build_typeshed_third_party_package_map(&typeshed_path); + let first_name_part = module_descriptor.name_parts.first().map(String::as_str)?; + package_paths.get(first_name_part).cloned() +}