Skip to content

Commit

Permalink
[red-knot] Add support for editable installs to the module resolver (#…
Browse files Browse the repository at this point in the history
…12307)

Co-authored-by: Micha Reiser <micha@reiser.io>
Co-authored-by: Carl Meyer <carl@astral.sh>
  • Loading branch information
3 people committed Jul 16, 2024
1 parent 595b1aa commit 9a2dafb
Show file tree
Hide file tree
Showing 4 changed files with 687 additions and 41 deletions.
3 changes: 2 additions & 1 deletion crates/red_knot_module_resolver/src/db.rs
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
use ruff_db::Upcast;

use crate::resolver::{
file_to_module,
editable_install_resolution_paths, file_to_module,
internal::{ModuleNameIngredient, ModuleResolverSettings},
resolve_module_query,
};
Expand All @@ -11,6 +11,7 @@ use crate::typeshed::parse_typeshed_versions;
pub struct Jar(
ModuleNameIngredient<'_>,
ModuleResolverSettings,
editable_install_resolution_paths,
resolve_module_query,
file_to_module,
parse_typeshed_versions,
Expand Down
77 changes: 75 additions & 2 deletions crates/red_knot_module_resolver/src/path.rs
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@
use std::fmt;

use ruff_db::files::{system_path_to_file, vendored_path_to_file, File, FilePath};
use ruff_db::system::{SystemPath, SystemPathBuf};
use ruff_db::system::{System, SystemPath, SystemPathBuf};
use ruff_db::vendored::{VendoredPath, VendoredPathBuf};

use crate::db::Db;
Expand Down Expand Up @@ -73,6 +73,7 @@ enum ModuleResolutionPathBufInner {
FirstParty(SystemPathBuf),
StandardLibrary(FilePath),
SitePackages(SystemPathBuf),
EditableInstall(SystemPathBuf),
}

impl ModuleResolutionPathBufInner {
Expand Down Expand Up @@ -134,6 +135,19 @@ impl ModuleResolutionPathBufInner {
);
path.push(component);
}
Self::EditableInstall(ref mut path) => {
if let Some(extension) = extension {
assert!(
matches!(extension, "pyi" | "py"),
"Extension must be `py` or `pyi`; got `{extension}`"
);
}
assert!(
path.extension().is_none(),
"Cannot push part {component} to {path}, which already has an extension"
);
path.push(component);
}
}
}
}
Expand Down Expand Up @@ -197,6 +211,18 @@ impl ModuleResolutionPathBuf {
.then_some(Self(ModuleResolutionPathBufInner::SitePackages(path)))
}

#[must_use]
pub(crate) fn editable_installation_root(
system: &dyn System,
path: impl Into<SystemPathBuf>,
) -> Option<Self> {
let path = path.into();
// TODO: Add Salsa invalidation to this system call:
system
.is_directory(&path)
.then_some(Self(ModuleResolutionPathBufInner::EditableInstall(path)))
}

#[must_use]
pub(crate) fn is_regular_package(&self, search_path: &Self, resolver: &ResolverState) -> bool {
ModuleResolutionPathRef::from(self).is_regular_package(search_path, resolver)
Expand Down Expand Up @@ -229,6 +255,16 @@ impl ModuleResolutionPathBuf {
pub(crate) fn to_file(&self, search_path: &Self, resolver: &ResolverState) -> Option<File> {
ModuleResolutionPathRef::from(self).to_file(search_path, resolver)
}

pub(crate) fn as_system_path(&self) -> Option<&SystemPathBuf> {
match &self.0 {
ModuleResolutionPathBufInner::Extra(path) => Some(path),
ModuleResolutionPathBufInner::FirstParty(path) => Some(path),
ModuleResolutionPathBufInner::StandardLibrary(_) => None,
ModuleResolutionPathBufInner::SitePackages(path) => Some(path),
ModuleResolutionPathBufInner::EditableInstall(path) => Some(path),
}
}
}

impl fmt::Debug for ModuleResolutionPathBuf {
Expand All @@ -250,6 +286,10 @@ impl fmt::Debug for ModuleResolutionPathBuf {
.debug_tuple("ModuleResolutionPathBuf::StandardLibrary")
.field(path)
.finish(),
ModuleResolutionPathBufInner::EditableInstall(path) => f
.debug_tuple("ModuleResolutionPathBuf::EditableInstall")
.field(path)
.finish(),
}
}
}
Expand All @@ -272,6 +312,7 @@ enum ModuleResolutionPathRefInner<'a> {
FirstParty(&'a SystemPath),
StandardLibrary(FilePathRef<'a>),
SitePackages(&'a SystemPath),
EditableInstall(&'a SystemPath),
}

impl<'a> ModuleResolutionPathRefInner<'a> {
Expand Down Expand Up @@ -306,6 +347,7 @@ impl<'a> ModuleResolutionPathRefInner<'a> {
(Self::Extra(path), Self::Extra(_)) => resolver.system().is_directory(path),
(Self::FirstParty(path), Self::FirstParty(_)) => resolver.system().is_directory(path),
(Self::SitePackages(path), Self::SitePackages(_)) => resolver.system().is_directory(path),
(Self::EditableInstall(path), Self::EditableInstall(_)) => resolver.system().is_directory(path),
(Self::StandardLibrary(path), Self::StandardLibrary(stdlib_root)) => {
match Self::query_stdlib_version(path, search_path, &stdlib_root, resolver) {
TypeshedVersionsQueryResult::DoesNotExist => false,
Expand Down Expand Up @@ -333,6 +375,7 @@ impl<'a> ModuleResolutionPathRefInner<'a> {
(Self::Extra(path), Self::Extra(_)) => is_non_stdlib_pkg(resolver, path),
(Self::FirstParty(path), Self::FirstParty(_)) => is_non_stdlib_pkg(resolver, path),
(Self::SitePackages(path), Self::SitePackages(_)) => is_non_stdlib_pkg(resolver, path),
(Self::EditableInstall(path), Self::EditableInstall(_)) => is_non_stdlib_pkg(resolver, path),
// Unlike the other variants:
// (1) Account for VERSIONS
// (2) Only test for `__init__.pyi`, not `__init__.py`
Expand All @@ -358,6 +401,7 @@ impl<'a> ModuleResolutionPathRefInner<'a> {
(Self::SitePackages(path), Self::SitePackages(_)) => {
system_path_to_file(resolver.db.upcast(), path)
}
(Self::EditableInstall(path), Self::EditableInstall(_)) => system_path_to_file(resolver.db.upcast(), path),
(Self::StandardLibrary(path), Self::StandardLibrary(stdlib_root)) => {
match Self::query_stdlib_version(&path, search_path, &stdlib_root, resolver) {
TypeshedVersionsQueryResult::DoesNotExist => None,
Expand All @@ -374,7 +418,10 @@ impl<'a> ModuleResolutionPathRefInner<'a> {
#[must_use]
fn to_module_name(self) -> Option<ModuleName> {
match self {
Self::Extra(path) | Self::FirstParty(path) | Self::SitePackages(path) => {
Self::Extra(path)
| Self::FirstParty(path)
| Self::SitePackages(path)
| Self::EditableInstall(path) => {
let parent = path.parent()?;
let parent_components = parent.components().map(|component| component.as_str());
let skip_final_part =
Expand Down Expand Up @@ -421,6 +468,9 @@ impl<'a> ModuleResolutionPathRefInner<'a> {
Self::SitePackages(path) => {
ModuleResolutionPathBufInner::SitePackages(path.with_extension("pyi"))
}
Self::EditableInstall(path) => {
ModuleResolutionPathBufInner::EditableInstall(path.with_extension("pyi"))
}
}
}

Expand All @@ -437,6 +487,9 @@ impl<'a> ModuleResolutionPathRefInner<'a> {
Self::SitePackages(path) => Some(ModuleResolutionPathBufInner::SitePackages(
path.with_extension("py"),
)),
Self::EditableInstall(path) => Some(ModuleResolutionPathBufInner::EditableInstall(
path.with_extension("py"),
)),
}
}

Expand Down Expand Up @@ -474,6 +527,13 @@ impl<'a> ModuleResolutionPathRefInner<'a> {
.then_some(Self::SitePackages(path))
})
}
(Self::EditableInstall(root), FilePathRef::System(absolute_path)) => {
absolute_path.strip_prefix(root).ok().and_then(|path| {
path.extension()
.map_or(true, |ext| matches!(ext, "pyi" | "py"))
.then_some(Self::EditableInstall(path))
})
}
(Self::Extra(_), FilePathRef::Vendored(_)) => None,
(Self::FirstParty(_), FilePathRef::Vendored(_)) => None,
(Self::StandardLibrary(root), FilePathRef::Vendored(absolute_path)) => match root {
Expand All @@ -487,6 +547,7 @@ impl<'a> ModuleResolutionPathRefInner<'a> {
}
},
(Self::SitePackages(_), FilePathRef::Vendored(_)) => None,
(Self::EditableInstall(_), FilePathRef::Vendored(_)) => None,
}
}
}
Expand Down Expand Up @@ -562,6 +623,10 @@ impl fmt::Debug for ModuleResolutionPathRef<'_> {
.debug_tuple("ModuleResolutionPathRef::StandardLibrary")
.field(path)
.finish(),
ModuleResolutionPathRefInner::EditableInstall(path) => f
.debug_tuple("ModuleResolutionPathRef::EditableInstall")
.field(path)
.finish(),
}
}
}
Expand All @@ -582,6 +647,9 @@ impl<'a> From<&'a ModuleResolutionPathBuf> for ModuleResolutionPathRef<'a> {
ModuleResolutionPathBufInner::SitePackages(path) => {
ModuleResolutionPathRefInner::SitePackages(path)
}
ModuleResolutionPathBufInner::EditableInstall(path) => {
ModuleResolutionPathRefInner::EditableInstall(path)
}
};
ModuleResolutionPathRef(inner)
}
Expand All @@ -593,6 +661,7 @@ impl PartialEq<SystemPath> for ModuleResolutionPathRef<'_> {
ModuleResolutionPathRefInner::Extra(path) => path == other,
ModuleResolutionPathRefInner::FirstParty(path) => path == other,
ModuleResolutionPathRefInner::SitePackages(path) => path == other,
ModuleResolutionPathRefInner::EditableInstall(path) => path == other,
ModuleResolutionPathRefInner::StandardLibrary(FilePathRef::System(path)) => {
path == other
}
Expand Down Expand Up @@ -625,6 +694,7 @@ impl PartialEq<VendoredPath> for ModuleResolutionPathRef<'_> {
ModuleResolutionPathRefInner::Extra(_) => false,
ModuleResolutionPathRefInner::FirstParty(_) => false,
ModuleResolutionPathRefInner::SitePackages(_) => false,
ModuleResolutionPathRefInner::EditableInstall(_) => false,
ModuleResolutionPathRefInner::StandardLibrary(FilePathRef::System(_)) => false,
ModuleResolutionPathRefInner::StandardLibrary(FilePathRef::Vendored(path)) => {
path == other
Expand Down Expand Up @@ -707,6 +777,9 @@ mod tests {
ModuleResolutionPathRefInner::SitePackages(path) => {
ModuleResolutionPathBufInner::SitePackages(path.to_path_buf())
}
ModuleResolutionPathRefInner::EditableInstall(path) => {
ModuleResolutionPathBufInner::EditableInstall(path.to_path_buf())
}
};
ModuleResolutionPathBuf(inner)
}
Expand Down
Loading

0 comments on commit 9a2dafb

Please sign in to comment.