From 1ed227a1e0dc18a0e12f90d10757a89883a9eb01 Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Tue, 27 Jun 2023 12:15:07 -0400 Subject: [PATCH] Port Pyright's import resolver to Rust (#5381) ## Summary This PR contains the first step towards enabling robust first-party, third-party, and standard library import resolution in Ruff (including support for `typeshed`, stub files, native modules, etc.) by porting Pyright's import resolver to Rust. The strategy taken here was to start with a more-or-less direct port of the Pyright's TypeScript resolver. The code is intentionally similar, and the test suite is effectively a superset of Pyright's test suite for its own resolver. Due to the nature of the port, the code is very, very non-idiomatic for Rust. The code is also entirely unused outside of the test suite, and no effort has been made to integrate it with the rest of the codebase. Future work will include: - Refactoring the code (now that it works) to match Rust and Ruff idioms. - Further testing, in practice, to ensure that the resolver can resolve imports in a complex project, when provided with a virtual environment path. - Caching, to minimize filesystem lookups and redundant resolutions. - Integration into Ruff itself (use Ruff's existing settings, find rules that can make use of robust resolution, etc.) --- Cargo.lock | 37 + LICENSE | 26 + README.md | 4 +- crates/ruff_python_resolver/Cargo.toml | 21 + crates/ruff_python_resolver/src/config.rs | 26 + .../src/execution_environment.rs | 19 + crates/ruff_python_resolver/src/host.rs | 43 + .../src/implicit_imports.rs | 150 ++ .../ruff_python_resolver/src/import_result.rs | 122 ++ crates/ruff_python_resolver/src/lib.rs | 16 + .../src/module_descriptor.rs | 16 + .../ruff_python_resolver/src/native_module.rs | 14 + crates/ruff_python_resolver/src/py_typed.rs | 40 + .../src/python_platform.rs | 7 + .../src/python_version.rs | 24 + crates/ruff_python_resolver/src/resolver.rs | 1497 +++++++++++++++++ crates/ruff_python_resolver/src/search.rs | 282 ++++ 17 files changed, 2343 insertions(+), 1 deletion(-) create mode 100644 crates/ruff_python_resolver/Cargo.toml create mode 100644 crates/ruff_python_resolver/src/config.rs create mode 100644 crates/ruff_python_resolver/src/execution_environment.rs create mode 100644 crates/ruff_python_resolver/src/host.rs create mode 100644 crates/ruff_python_resolver/src/implicit_imports.rs create mode 100644 crates/ruff_python_resolver/src/import_result.rs create mode 100644 crates/ruff_python_resolver/src/lib.rs create mode 100644 crates/ruff_python_resolver/src/module_descriptor.rs create mode 100644 crates/ruff_python_resolver/src/native_module.rs create mode 100644 crates/ruff_python_resolver/src/py_typed.rs create mode 100644 crates/ruff_python_resolver/src/python_platform.rs create mode 100644 crates/ruff_python_resolver/src/python_version.rs create mode 100644 crates/ruff_python_resolver/src/resolver.rs create mode 100644 crates/ruff_python_resolver/src/search.rs 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() +}