This commit is contained in:
Charlie Marsh 2023-07-25 22:53:29 -04:00
parent 15273c6d95
commit 68a12f72a4
2 changed files with 155 additions and 24 deletions

View File

@ -4,6 +4,142 @@ use smallvec::{smallvec, SmallVec};
/// 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<(&ast::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.

View File

@ -2,9 +2,11 @@ use std::path::Path;
use bitflags::bitflags;
use rustc_hash::FxHashMap;
use smallvec::{smallvec, SmallVec};
use smallvec::smallvec;
use ruff_python_ast::call_path::{collect_call_path, from_unqualified_name, CallPath};
use ruff_python_ast::call_path::{
collect_call_path, collect_head_path, from_unqualified_name, CallPath,
};
use ruff_python_ast::{self as ast, Expr, Ranged, Stmt};
use ruff_python_stdlib::path::is_python_stub_file;
use ruff_python_stdlib::typing::is_typing_extension;
@ -623,40 +625,33 @@ impl<'a> SemanticModel<'a> {
}
}
let (head, tail) = collect_head_path(value)?;
// If the name was already resolved, look it up; otherwise, search for the symbol.
let head = match_head(value)?;
let binding = self
.resolved_names
.get(&head.into())
.map(|id| self.binding(*id))
.or_else(|| self.find_binding(&head.id))?;
let binding = if let Some(id) = self.resolved_names.get(&head.into()) {
self.binding(*id)
} else {
self.find_binding(&head.id)?
};
match &binding.kind {
BindingKind::Import(Import { call_path, .. }) => {
let x = collect_call_path(value)?;
let (_, tail) = x.split_first()?;
let mut resolved = SmallVec::with_capacity(call_path.len() + tail.len());
resolved.extend_from_slice(call_path);
resolved.extend_from_slice(tail);
let resolved: CallPath = call_path.iter().chain(tail.iter()).copied().collect();
// resolved.extend_from_slice(call_path);
// resolved.extend_from_slice(tail);
Some(resolved)
}
BindingKind::SubmoduleImport(SubmoduleImport { qualified_name, .. }) => {
let x = collect_call_path(value)?;
let (_, tail) = x.split_first()?;
let name = qualified_name.split('.').next().unwrap_or(qualified_name);
let mut source_path: CallPath = from_unqualified_name(name);
source_path.extend_from_slice(tail);
source_path.extend_from_slice(tail.as_slice());
Some(source_path)
}
BindingKind::FromImport(FromImport { call_path, .. }) => {
let x = collect_call_path(value)?;
let (_, tail) = x.split_first()?;
let mut resolved = SmallVec::with_capacity(call_path.len() + tail.len());
resolved.extend_from_slice(call_path);
resolved.extend_from_slice(tail);
let resolved: CallPath = call_path.iter().chain(tail.iter()).copied().collect();
// let mut resolved = SmallVec::with_capacity(call_path.len() + tail.len());
// resolved.extend_from_slice(call_path);
// resolved.extend_from_slice(tail);
Some(resolved)
}
BindingKind::Builtin => Some(smallvec!["", head.id.as_str()]),