This commit is contained in:
Charlie Marsh
2023-07-27 21:08:08 -04:00
parent ac0f51ef4f
commit ac25bd9596
13 changed files with 202 additions and 339 deletions

View File

@@ -5,142 +5,6 @@ use crate::{nodes, Expr};
/// A representation of a qualified name, like `typing.List`.
pub type CallPath<'a> = SmallVec<[&'a str; 8]>;
/// Convert an `Expr` to its [`CallPath`] segments (like `["typing", "List"]`).
pub fn collect_head_path(expr: &Expr) -> Option<(&nodes::ExprName, CallPath)> {
// Unroll the loop up to eight times, to match the maximum number of expected attributes.
// In practice, unrolling appears to give about a 4x speed-up on this hot path.
let attr1 = match expr {
Expr::Attribute(attr1) => attr1,
// Ex) `foo`
Expr::Name(name) => return Some((name, CallPath::new())),
_ => return None,
};
let attr2 = match attr1.value.as_ref() {
Expr::Attribute(attr2) => attr2,
// Ex) `foo.bar`
Expr::Name(name) => {
return Some((name, CallPath::from_slice(&[attr1.attr.as_str()])));
}
_ => return None,
};
let attr3 = match attr2.value.as_ref() {
Expr::Attribute(attr3) => attr3,
// Ex) `foo.bar.baz`
Expr::Name(name) => {
return Some((
name,
CallPath::from_slice(&[attr2.attr.as_str(), attr1.attr.as_str()]),
));
}
_ => return None,
};
let attr4 = match attr3.value.as_ref() {
Expr::Attribute(attr4) => attr4,
// Ex) `foo.bar.baz.bop`
Expr::Name(name) => {
return Some((
name,
CallPath::from_slice(&[
attr3.attr.as_str(),
attr2.attr.as_str(),
attr1.attr.as_str(),
]),
));
}
_ => return None,
};
let attr5 = match attr4.value.as_ref() {
Expr::Attribute(attr5) => attr5,
// Ex) `foo.bar.baz.bop.bap`
Expr::Name(name) => {
return Some((
name,
CallPath::from_slice(&[
attr4.attr.as_str(),
attr3.attr.as_str(),
attr2.attr.as_str(),
attr1.attr.as_str(),
]),
));
}
_ => return None,
};
let attr6 = match attr5.value.as_ref() {
Expr::Attribute(attr6) => attr6,
// Ex) `foo.bar.baz.bop.bap.bab`
Expr::Name(name) => {
return Some((
name,
CallPath::from_slice(&[
attr5.attr.as_str(),
attr4.attr.as_str(),
attr3.attr.as_str(),
attr2.attr.as_str(),
attr1.attr.as_str(),
]),
));
}
_ => return None,
};
let attr7 = match attr6.value.as_ref() {
Expr::Attribute(attr7) => attr7,
// Ex) `foo.bar.baz.bop.bap.bab.bob`
Expr::Name(name) => {
return Some((
name,
CallPath::from_slice(&[
attr6.attr.as_str(),
attr5.attr.as_str(),
attr4.attr.as_str(),
attr3.attr.as_str(),
attr2.attr.as_str(),
attr1.attr.as_str(),
]),
));
}
_ => return None,
};
let attr8 = match attr7.value.as_ref() {
Expr::Attribute(attr8) => attr8,
// Ex) `foo.bar.baz.bop.bap.bab.bob.bib`
Expr::Name(name) => {
return Some((
name,
CallPath::from_slice(&[
attr7.attr.as_str(),
attr6.attr.as_str(),
attr5.attr.as_str(),
attr4.attr.as_str(),
attr3.attr.as_str(),
attr2.attr.as_str(),
attr1.attr.as_str(),
]),
));
}
_ => return None,
};
let (name, mut call_path) = collect_head_path(&attr8.value)?;
call_path.extend([
attr8.attr.as_str(),
attr7.attr.as_str(),
attr6.attr.as_str(),
attr5.attr.as_str(),
attr4.attr.as_str(),
attr3.attr.as_str(),
attr2.attr.as_str(),
attr1.attr.as_str(),
]);
Some((name, call_path))
}
/// Convert an `Expr` to its [`CallPath`] segments (like `["typing", "List"]`).
pub fn collect_call_path(expr: &Expr) -> Option<CallPath> {
// Unroll the loop up to eight times, to match the maximum number of expected attributes.
@@ -291,7 +155,7 @@ pub fn format_call_path(call_path: &[&str]) -> String {
let mut formatted = String::new();
let mut iter = call_path.iter();
for segment in iter.by_ref() {
if matches!(*segment, ".") {
if *segment == "." {
formatted.push('.');
} else {
formatted.push_str(segment);

View File

@@ -856,78 +856,18 @@ pub fn to_module_path(package: &Path, path: &Path) -> Option<Vec<String>> {
.collect::<Option<Vec<String>>>()
}
/// Create a [`CallPath`] from a relative import reference name (like `".foo.bar"`).
///
/// Returns an empty [`CallPath`] if the import is invalid (e.g., a relative import that
/// extends beyond the top-level module).
/// Format the call path for a relative import.
///
/// # Examples
///
/// ```rust
/// # use smallvec::{smallvec, SmallVec};
/// # use ruff_python_ast::helpers::from_relative_import;
/// # use ruff_python_ast::helpers::collect_import_from_member;
///
/// assert_eq!(from_relative_import(&[], "bar"), SmallVec::from_buf(["bar"]));
/// assert_eq!(from_relative_import(&["foo".to_string()], "bar"), SmallVec::from_buf(["foo", "bar"]));
/// assert_eq!(from_relative_import(&["foo".to_string()], "bar.baz"), SmallVec::from_buf(["foo", "bar", "baz"]));
/// assert_eq!(from_relative_import(&["foo".to_string()], ".bar"), SmallVec::from_buf(["bar"]));
/// assert!(from_relative_import(&["foo".to_string()], "..bar").is_empty());
/// assert!(from_relative_import(&["foo".to_string()], "...bar").is_empty());
/// assert_eq!(collect_import_from_member(None, None, "bar").as_slice(), ["bar"]);
/// assert_eq!(collect_import_from_member(Some(1), None, "bar").as_slice(), [".", "bar"]);
/// assert_eq!(collect_import_from_member(Some(1), Some("foo"), "bar").as_slice(), [".", "foo", "bar"]);
/// ```
pub fn from_relative_import<'a>(module: &'a [String], name: &'a str) -> CallPath<'a> {
let mut call_path: CallPath = SmallVec::with_capacity(module.len() + 1);
// Start with the module path.
call_path.extend(module.iter().map(String::as_str));
// Remove segments based on the number of dots.
for _ in 0..name.chars().take_while(|c| *c == '.').count() {
if call_path.is_empty() {
return SmallVec::new();
}
call_path.pop();
}
// Add the remaining segments.
call_path.extend(name.trim_start_matches('.').split('.'));
call_path
}
pub fn from_relative_import_parts<'a>(
module_path: &'a [String],
level: Option<u32>,
module: Option<&'a str>,
member: &'a str,
) -> Option<CallPath<'a>> {
let mut call_path: CallPath = SmallVec::with_capacity(module_path.len() + 1);
// Remove segments based on the number of dots.
if let Some(level) = level {
if level > 0 {
call_path.extend(module_path.iter().map(String::as_str));
for _ in 0..level {
if call_path.is_empty() {
break;
}
call_path.pop();
}
}
}
// Add the remaining segments.
if let Some(module) = module {
call_path.extend(module.split('.'));
}
// Add the member.
call_path.push(member);
Some(call_path)
}
pub fn literal_path<'a>(
pub fn collect_import_from_member<'a>(
level: Option<u32>,
module: Option<&'a str>,
member: &'a str,
@@ -940,7 +880,7 @@ pub fn literal_path<'a>(
+ 1,
);
// Include the dots
// Include the dots as standalone segments.
if let Some(level) = level {
if level > 0 {
for _ in 0..level {
@@ -960,6 +900,39 @@ pub fn literal_path<'a>(
call_path
}
/// Format the call path for a relative import, or `None` if the relative import extends beyond
/// the root module.
pub fn from_relative_import<'a>(
// The path from which the import is relative.
module: &'a [String],
// The path of the import itself (e.g., given `from ..foo import bar`, `[".", ".", "foo", "bar]`).
import: &[&'a str],
// The remaining segments to the call path (e.g., given `bar.baz`, `["baz"]`).
tail: &[&'a str],
) -> Option<CallPath<'a>> {
let mut call_path: CallPath = SmallVec::with_capacity(module.len() + import.len() + tail.len());
// Start with the module path.
call_path.extend(module.iter().map(String::as_str));
// Remove segments based on the number of dots.
for segment in import {
if *segment == "." {
if call_path.is_empty() {
return None;
}
call_path.pop();
} else {
call_path.push(segment);
}
}
// Add the remaining segments.
call_path.extend_from_slice(tail);
Some(call_path)
}
/// Given an imported module (based on its relative import level and module name), return the
/// fully-qualified module path.
pub fn resolve_imported_module_path<'a>(