Avoid unnecessary allocations for module names (#1863)

This commit is contained in:
Charlie Marsh 2023-01-13 23:26:34 -05:00 committed by GitHub
parent 7b1ce72f86
commit 59542344e2
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 122 additions and 69 deletions

View File

@ -104,7 +104,7 @@ impl<'a> Scope<'a> {
} }
#[derive(Clone, Debug)] #[derive(Clone, Debug)]
pub enum BindingKind { pub enum BindingKind<'a> {
Annotation, Annotation,
Argument, Argument,
Assignment, Assignment,
@ -118,14 +118,14 @@ pub enum BindingKind {
Export(Vec<String>), Export(Vec<String>),
FutureImportation, FutureImportation,
StarImportation(Option<usize>, Option<String>), StarImportation(Option<usize>, Option<String>),
Importation(String, String), Importation(&'a str, &'a str),
FromImportation(String, String), FromImportation(&'a str, String),
SubmoduleImportation(String, String), SubmoduleImportation(&'a str, &'a str),
} }
#[derive(Clone, Debug)] #[derive(Clone, Debug)]
pub struct Binding<'a> { pub struct Binding<'a> {
pub kind: BindingKind, pub kind: BindingKind<'a>,
pub range: Range, pub range: Range,
/// The statement in which the `Binding` was defined. /// The statement in which the `Binding` was defined.
pub source: Option<RefEquality<'a, Stmt>>, pub source: Option<RefEquality<'a, Stmt>>,
@ -168,19 +168,26 @@ impl<'a> Binding<'a> {
pub fn redefines(&self, existing: &'a Binding) -> bool { pub fn redefines(&self, existing: &'a Binding) -> bool {
match &self.kind { match &self.kind {
BindingKind::Importation(_, full_name) | BindingKind::FromImportation(_, full_name) => { BindingKind::Importation(.., full_name) => {
if let BindingKind::SubmoduleImportation(_, existing_full_name) = &existing.kind { if let BindingKind::SubmoduleImportation(.., existing) = &existing.kind {
return full_name == existing_full_name; return full_name == existing;
} }
} }
BindingKind::SubmoduleImportation(_, full_name) => { BindingKind::FromImportation(.., full_name) => {
if let BindingKind::Importation(_, existing_full_name) if let BindingKind::SubmoduleImportation(.., existing) = &existing.kind {
| BindingKind::FromImportation(_, existing_full_name) return full_name == existing;
| BindingKind::SubmoduleImportation(_, existing_full_name) = &existing.kind
{
return full_name == existing_full_name;
} }
} }
BindingKind::SubmoduleImportation(.., full_name) => match &existing.kind {
BindingKind::Importation(.., existing)
| BindingKind::SubmoduleImportation(.., existing) => {
return full_name == existing;
}
BindingKind::FromImportation(.., existing) => {
return full_name == existing;
}
_ => {}
},
BindingKind::Annotation => { BindingKind::Annotation => {
return false; return false;
} }

View File

@ -209,22 +209,41 @@ impl<'a> Checker<'a> {
let call_path = collect_call_path(value); let call_path = collect_call_path(value);
if let Some(head) = call_path.first() { if let Some(head) = call_path.first() {
if let Some(binding) = self.find_binding(head) { if let Some(binding) = self.find_binding(head) {
if let BindingKind::Importation(.., name) match &binding.kind {
| BindingKind::SubmoduleImportation(name, ..) BindingKind::Importation(.., name) => {
| BindingKind::FromImportation(.., name) = &binding.kind // Ignore relative imports.
{ if name.starts_with('.') {
// Ignore relative imports. return None;
if name.starts_with('.') { }
return None; let mut source_path: Vec<&str> = name.split('.').collect();
source_path.extend(call_path.iter().skip(1));
return Some(source_path);
} }
let mut source_path: Vec<&str> = name.split('.').collect(); BindingKind::SubmoduleImportation(name, ..) => {
source_path.extend(call_path.iter().skip(1)); // Ignore relative imports.
return Some(source_path); if name.starts_with('.') {
} else if let BindingKind::Builtin = &binding.kind { return None;
let mut source_path: Vec<&str> = Vec::with_capacity(call_path.len() + 1); }
source_path.push(""); let mut source_path: Vec<&str> = name.split('.').collect();
source_path.extend(call_path); source_path.extend(call_path.iter().skip(1));
return Some(source_path); return Some(source_path);
}
BindingKind::FromImportation(.., name) => {
// Ignore relative imports.
if name.starts_with('.') {
return None;
}
let mut source_path: Vec<&str> = name.split('.').collect();
source_path.extend(call_path.iter().skip(1));
return Some(source_path);
}
BindingKind::Builtin => {
let mut source_path: Vec<&str> = Vec::with_capacity(call_path.len() + 1);
source_path.push("");
source_path.extend(call_path);
return Some(source_path);
}
_ => {}
} }
} }
} }
@ -724,10 +743,7 @@ where
self.add_binding( self.add_binding(
name, name,
Binding { Binding {
kind: BindingKind::SubmoduleImportation( kind: BindingKind::SubmoduleImportation(name, full_name),
name.to_string(),
full_name.to_string(),
),
used: None, used: None,
range: Range::from_located(alias), range: Range::from_located(alias),
source: Some(self.current_stmt().clone()), source: Some(self.current_stmt().clone()),
@ -746,10 +762,7 @@ where
self.add_binding( self.add_binding(
name, name,
Binding { Binding {
kind: BindingKind::Importation( kind: BindingKind::Importation(name, full_name),
name.to_string(),
full_name.to_string(),
),
// Treat explicit re-export as usage (e.g., `import applications // Treat explicit re-export as usage (e.g., `import applications
// as applications`). // as applications`).
used: if alias used: if alias
@ -1062,7 +1075,7 @@ where
self.add_binding( self.add_binding(
name, name,
Binding { Binding {
kind: BindingKind::FromImportation(name.to_string(), full_name), kind: BindingKind::FromImportation(name, full_name),
// Treat explicit re-export as usage (e.g., `from .applications // Treat explicit re-export as usage (e.g., `from .applications
// import FastAPI as FastAPI`). // import FastAPI as FastAPI`).
used: if alias used: if alias
@ -2118,7 +2131,7 @@ where
if let BindingKind::Importation(.., module) = if let BindingKind::Importation(.., module) =
&binding.kind &binding.kind
{ {
module != "pandas" module != &"pandas"
} else { } else {
matches!( matches!(
binding.kind, binding.kind,
@ -3359,23 +3372,37 @@ impl<'a> Checker<'a> {
// import pyarrow as pa // import pyarrow as pa
// import pyarrow.csv // import pyarrow.csv
// print(pa.csv.read_csv("test.csv")) // print(pa.csv.read_csv("test.csv"))
if let BindingKind::Importation(name, full_name) match &self.bindings[*index].kind {
| BindingKind::FromImportation(name, full_name) BindingKind::Importation(name, full_name)
| BindingKind::SubmoduleImportation(name, full_name) = | BindingKind::SubmoduleImportation(name, full_name) => {
&self.bindings[*index].kind let has_alias = full_name
{ .split('.')
let has_alias = full_name .last()
.split('.') .map(|segment| &segment != name)
.last() .unwrap_or_default();
.map(|segment| segment != name) if has_alias {
.unwrap_or_default(); // Mark the sub-importation as used.
if has_alias { if let Some(index) = scope.values.get(full_name) {
// Mark the sub-importation as used. self.bindings[*index].used =
if let Some(index) = scope.values.get(full_name.as_str()) { Some((scope_id, Range::from_located(expr)));
self.bindings[*index].used = }
Some((scope_id, Range::from_located(expr)));
} }
} }
BindingKind::FromImportation(name, full_name) => {
let has_alias = full_name
.split('.')
.last()
.map(|segment| &segment != name)
.unwrap_or_default();
if has_alias {
// Mark the sub-importation as used.
if let Some(index) = scope.values.get(full_name.as_str()) {
self.bindings[*index].used =
Some((scope_id, Range::from_located(expr)));
}
}
}
_ => {}
} }
return; return;
@ -3901,9 +3928,12 @@ impl<'a> Checker<'a> {
{ {
let binding = &self.bindings[*index]; let binding = &self.bindings[*index];
let (BindingKind::Importation(_, full_name) let full_name = match &binding.kind {
| BindingKind::SubmoduleImportation(_, full_name) BindingKind::Importation(.., full_name) => full_name,
| BindingKind::FromImportation(_, full_name)) = &binding.kind else { continue; }; BindingKind::FromImportation(.., full_name) => full_name.as_str(),
BindingKind::SubmoduleImportation(.., full_name) => full_name,
_ => continue,
};
// Skip used exports from `__all__` // Skip used exports from `__all__`
if binding.used.is_some() if binding.used.is_some()

View File

@ -31,28 +31,44 @@ fn get_member_import_name_alias(checker: &Checker, module: &str, member: &str) -
// e.g. module=sys object=exit // e.g. module=sys object=exit
// `import sys` -> `sys.exit` // `import sys` -> `sys.exit`
// `import sys as sys2` -> `sys2.exit` // `import sys as sys2` -> `sys2.exit`
BindingKind::Importation(name, full_name) if full_name == module => { BindingKind::Importation(name, full_name) => {
Some(format!("{name}.{member}")) if full_name == &module {
Some(format!("{name}.{member}"))
} else {
None
}
} }
// e.g. module=os.path object=join // e.g. module=os.path object=join
// `from os.path import join` -> `join` // `from os.path import join` -> `join`
// `from os.path import join as join2` -> `join2` // `from os.path import join as join2` -> `join2`
BindingKind::FromImportation(name, full_name) BindingKind::FromImportation(name, full_name) => {
if full_name == &format!("{module}.{member}") => let mut parts = full_name.split('.');
{ if parts.next() == Some(module)
Some(name.to_string()) && parts.next() == Some(member)
&& parts.next().is_none()
{
Some((*name).to_string())
} else {
None
}
} }
// e.g. module=os.path object=join // e.g. module=os.path object=join
// `from os.path import *` -> `join` // `from os.path import *` -> `join`
BindingKind::StarImportation(_, name) BindingKind::StarImportation(_, name) => {
if name.as_ref().map(|name| name == module).unwrap_or_default() => if name.as_ref().map(|name| name == module).unwrap_or_default() {
{ Some(member.to_string())
Some(member.to_string()) } else {
None
}
} }
// e.g. module=os.path object=join // e.g. module=os.path object=join
// `import os.path ` -> `os.path.join` // `import os.path ` -> `os.path.join`
BindingKind::SubmoduleImportation(_, full_name) if full_name == module => { BindingKind::SubmoduleImportation(_, full_name) => {
Some(format!("{full_name}.{member}")) if full_name == &module {
Some(format!("{full_name}.{member}"))
} else {
None
}
} }
// Non-imports. // Non-imports.
_ => None, _ => None,