From 59542344e2f8c5010c0f80cbf3f8d404efeb845c Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Fri, 13 Jan 2023 23:26:34 -0500 Subject: [PATCH] Avoid unnecessary allocations for module names (#1863) --- src/ast/types.rs | 35 ++++++---- src/checkers/ast.rs | 116 +++++++++++++++++++------------ src/pylint/rules/use_sys_exit.rs | 40 +++++++---- 3 files changed, 122 insertions(+), 69 deletions(-) diff --git a/src/ast/types.rs b/src/ast/types.rs index e631e4a31c..57c6f3c44e 100644 --- a/src/ast/types.rs +++ b/src/ast/types.rs @@ -104,7 +104,7 @@ impl<'a> Scope<'a> { } #[derive(Clone, Debug)] -pub enum BindingKind { +pub enum BindingKind<'a> { Annotation, Argument, Assignment, @@ -118,14 +118,14 @@ pub enum BindingKind { Export(Vec), FutureImportation, StarImportation(Option, Option), - Importation(String, String), - FromImportation(String, String), - SubmoduleImportation(String, String), + Importation(&'a str, &'a str), + FromImportation(&'a str, String), + SubmoduleImportation(&'a str, &'a str), } #[derive(Clone, Debug)] pub struct Binding<'a> { - pub kind: BindingKind, + pub kind: BindingKind<'a>, pub range: Range, /// The statement in which the `Binding` was defined. pub source: Option>, @@ -168,19 +168,26 @@ impl<'a> Binding<'a> { pub fn redefines(&self, existing: &'a Binding) -> bool { match &self.kind { - BindingKind::Importation(_, full_name) | BindingKind::FromImportation(_, full_name) => { - if let BindingKind::SubmoduleImportation(_, existing_full_name) = &existing.kind { - return full_name == existing_full_name; + BindingKind::Importation(.., full_name) => { + if let BindingKind::SubmoduleImportation(.., existing) = &existing.kind { + return full_name == existing; } } - BindingKind::SubmoduleImportation(_, full_name) => { - if let BindingKind::Importation(_, existing_full_name) - | BindingKind::FromImportation(_, existing_full_name) - | BindingKind::SubmoduleImportation(_, existing_full_name) = &existing.kind - { - return full_name == existing_full_name; + BindingKind::FromImportation(.., full_name) => { + if let BindingKind::SubmoduleImportation(.., existing) = &existing.kind { + return full_name == existing; } } + 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 => { return false; } diff --git a/src/checkers/ast.rs b/src/checkers/ast.rs index 77297add5b..4436f70bf0 100644 --- a/src/checkers/ast.rs +++ b/src/checkers/ast.rs @@ -209,22 +209,41 @@ impl<'a> Checker<'a> { let call_path = collect_call_path(value); if let Some(head) = call_path.first() { if let Some(binding) = self.find_binding(head) { - if let BindingKind::Importation(.., name) - | BindingKind::SubmoduleImportation(name, ..) - | BindingKind::FromImportation(.., name) = &binding.kind - { - // Ignore relative imports. - if name.starts_with('.') { - return None; + match &binding.kind { + BindingKind::Importation(.., 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); } - let mut source_path: Vec<&str> = name.split('.').collect(); - source_path.extend(call_path.iter().skip(1)); - return Some(source_path); - } else if let BindingKind::Builtin = &binding.kind { - 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); + BindingKind::SubmoduleImportation(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::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( name, Binding { - kind: BindingKind::SubmoduleImportation( - name.to_string(), - full_name.to_string(), - ), + kind: BindingKind::SubmoduleImportation(name, full_name), used: None, range: Range::from_located(alias), source: Some(self.current_stmt().clone()), @@ -746,10 +762,7 @@ where self.add_binding( name, Binding { - kind: BindingKind::Importation( - name.to_string(), - full_name.to_string(), - ), + kind: BindingKind::Importation(name, full_name), // Treat explicit re-export as usage (e.g., `import applications // as applications`). used: if alias @@ -1062,7 +1075,7 @@ where self.add_binding( name, 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 // import FastAPI as FastAPI`). used: if alias @@ -2118,7 +2131,7 @@ where if let BindingKind::Importation(.., module) = &binding.kind { - module != "pandas" + module != &"pandas" } else { matches!( binding.kind, @@ -3359,23 +3372,37 @@ impl<'a> Checker<'a> { // import pyarrow as pa // import pyarrow.csv // print(pa.csv.read_csv("test.csv")) - if let BindingKind::Importation(name, full_name) - | BindingKind::FromImportation(name, full_name) - | BindingKind::SubmoduleImportation(name, full_name) = - &self.bindings[*index].kind - { - 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))); + match &self.bindings[*index].kind { + BindingKind::Importation(name, full_name) + | BindingKind::SubmoduleImportation(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) { + 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; @@ -3901,9 +3928,12 @@ impl<'a> Checker<'a> { { let binding = &self.bindings[*index]; - let (BindingKind::Importation(_, full_name) - | BindingKind::SubmoduleImportation(_, full_name) - | BindingKind::FromImportation(_, full_name)) = &binding.kind else { continue; }; + let full_name = match &binding.kind { + BindingKind::Importation(.., full_name) => full_name, + BindingKind::FromImportation(.., full_name) => full_name.as_str(), + BindingKind::SubmoduleImportation(.., full_name) => full_name, + _ => continue, + }; // Skip used exports from `__all__` if binding.used.is_some() diff --git a/src/pylint/rules/use_sys_exit.rs b/src/pylint/rules/use_sys_exit.rs index 6a2bcdff1f..ccb5731a07 100644 --- a/src/pylint/rules/use_sys_exit.rs +++ b/src/pylint/rules/use_sys_exit.rs @@ -31,28 +31,44 @@ fn get_member_import_name_alias(checker: &Checker, module: &str, member: &str) - // e.g. module=sys object=exit // `import sys` -> `sys.exit` // `import sys as sys2` -> `sys2.exit` - BindingKind::Importation(name, full_name) if full_name == module => { - Some(format!("{name}.{member}")) + BindingKind::Importation(name, full_name) => { + if full_name == &module { + Some(format!("{name}.{member}")) + } else { + None + } } // e.g. module=os.path object=join // `from os.path import join` -> `join` // `from os.path import join as join2` -> `join2` - BindingKind::FromImportation(name, full_name) - if full_name == &format!("{module}.{member}") => - { - Some(name.to_string()) + BindingKind::FromImportation(name, full_name) => { + let mut parts = full_name.split('.'); + if parts.next() == Some(module) + && parts.next() == Some(member) + && parts.next().is_none() + { + Some((*name).to_string()) + } else { + None + } } // e.g. module=os.path object=join // `from os.path import *` -> `join` - BindingKind::StarImportation(_, name) - if name.as_ref().map(|name| name == module).unwrap_or_default() => - { - Some(member.to_string()) + BindingKind::StarImportation(_, name) => { + if name.as_ref().map(|name| name == module).unwrap_or_default() { + Some(member.to_string()) + } else { + None + } } // e.g. module=os.path object=join // `import os.path ` -> `os.path.join` - BindingKind::SubmoduleImportation(_, full_name) if full_name == module => { - Some(format!("{full_name}.{member}")) + BindingKind::SubmoduleImportation(_, full_name) => { + if full_name == &module { + Some(format!("{full_name}.{member}")) + } else { + None + } } // Non-imports. _ => None,