Improve performance of import matching code (#744)

This commit is contained in:
Charlie Marsh 2022-11-14 17:14:22 -05:00 committed by GitHub
parent fea029ae35
commit f67727b13c
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
13 changed files with 238 additions and 150 deletions

View file

@ -3,13 +3,14 @@ use once_cell::sync::Lazy;
use regex::Regex; use regex::Regex;
use rustpython_ast::{Excepthandler, ExcepthandlerKind, Expr, ExprKind, Location, StmtKind}; use rustpython_ast::{Excepthandler, ExcepthandlerKind, Expr, ExprKind, Location, StmtKind};
fn compose_call_path_inner<'a>(expr: &'a Expr, parts: &mut Vec<&'a str>) { #[inline(always)]
fn collect_call_path_inner<'a>(expr: &'a Expr, parts: &mut Vec<&'a str>) {
match &expr.node { match &expr.node {
ExprKind::Call { func, .. } => { ExprKind::Call { func, .. } => {
compose_call_path_inner(func, parts); collect_call_path_inner(func, parts);
} }
ExprKind::Attribute { value, attr, .. } => { ExprKind::Attribute { value, attr, .. } => {
compose_call_path_inner(value, parts); collect_call_path_inner(value, parts);
parts.push(attr); parts.push(attr);
} }
ExprKind::Name { id, .. } => { ExprKind::Name { id, .. } => {
@ -20,9 +21,9 @@ fn compose_call_path_inner<'a>(expr: &'a Expr, parts: &mut Vec<&'a str>) {
} }
/// Convert an `Expr` to its call path (like `List`, or `typing.List`). /// Convert an `Expr` to its call path (like `List`, or `typing.List`).
#[inline(always)]
pub fn compose_call_path(expr: &Expr) -> Option<String> { pub fn compose_call_path(expr: &Expr) -> Option<String> {
let mut segments = vec![]; let segments = collect_call_paths(expr);
compose_call_path_inner(expr, &mut segments);
if segments.is_empty() { if segments.is_empty() {
None None
} else { } else {
@ -30,6 +31,14 @@ pub fn compose_call_path(expr: &Expr) -> Option<String> {
} }
} }
/// Convert an `Expr` to its call path segments (like ["typing", "List"]).
#[inline(always)]
pub fn collect_call_paths(expr: &Expr) -> Vec<&str> {
let mut segments = vec![];
collect_call_path_inner(expr, &mut segments);
segments
}
/// Return `true` if the `Expr` is a name or attribute reference to `${target}`. /// Return `true` if the `Expr` is a name or attribute reference to `${target}`.
pub fn match_name_or_attr(expr: &Expr, target: &str) -> bool { pub fn match_name_or_attr(expr: &Expr, target: &str) -> bool {
match &expr.node { match &expr.node {
@ -72,7 +81,7 @@ pub fn is_assignment_to_a_dunder(node: &StmtKind) -> bool {
} }
/// Extract the names of all handled exceptions. /// Extract the names of all handled exceptions.
pub fn extract_handler_names(handlers: &[Excepthandler]) -> Vec<String> { pub fn extract_handler_names(handlers: &[Excepthandler]) -> Vec<Vec<&str>> {
let mut handler_names = vec![]; let mut handler_names = vec![];
for handler in handlers { for handler in handlers {
match &handler.node { match &handler.node {
@ -80,12 +89,16 @@ pub fn extract_handler_names(handlers: &[Excepthandler]) -> Vec<String> {
if let Some(type_) = type_ { if let Some(type_) = type_ {
if let ExprKind::Tuple { elts, .. } = &type_.node { if let ExprKind::Tuple { elts, .. } = &type_.node {
for type_ in elts { for type_ in elts {
if let Some(name) = compose_call_path(type_) { let call_path = collect_call_paths(type_);
handler_names.push(name); if !call_path.is_empty() {
handler_names.push(call_path);
} }
} }
} else if let Some(name) = compose_call_path(type_) { } else {
handler_names.push(name); let call_path = collect_call_paths(type_);
if !call_path.is_empty() {
handler_names.push(call_path);
}
} }
} }
} }
@ -123,66 +136,88 @@ pub fn to_absolute(relative: &Location, base: &Location) -> Location {
/// `typing.Union`. /// `typing.Union`.
pub fn match_module_member( pub fn match_module_member(
expr: &Expr, expr: &Expr,
target: &str, module: &str,
member: &str,
from_imports: &FnvHashMap<&str, FnvHashSet<&str>>, from_imports: &FnvHashMap<&str, FnvHashSet<&str>>,
) -> bool { ) -> bool {
compose_call_path(expr) match_call_path(&collect_call_paths(expr), module, member, from_imports)
.map(|expr| match_call_path(&expr, target, from_imports))
.unwrap_or(false)
} }
/// Return `true` if the `call_path` is a reference to `${module}.${target}`. /// Return `true` if the `call_path` is a reference to `${module}.${target}`.
/// ///
/// Optimized version of `match_module_member` for pre-computed call paths. /// Optimized version of `match_module_member` for pre-computed call paths.
pub fn match_call_path( pub fn match_call_path(
call_path: &str, call_path: &[&str],
target: &str, module: &str,
member: &str,
from_imports: &FnvHashMap<&str, FnvHashSet<&str>>, from_imports: &FnvHashMap<&str, FnvHashSet<&str>>,
) -> bool { ) -> bool {
// Case (1a): it's the same call path (`import typing`, `typing.re.Match`). // If we have no segments, we can't ever match.
// Case (1b): it's the same call path (`import typing.re`, `typing.re.Match`). let num_segments = call_path.len();
if call_path == target { if num_segments == 0 {
return false;
}
// If the last segment doesn't match the member, we can't ever match.
if call_path[num_segments - 1] != member {
return false;
}
// We now only need the module path, so throw out the member name.
let call_path = &call_path[..num_segments - 1];
let num_segments = call_path.len();
// Case (1a): We imported star from the parent (`from typing.re import *`,
// `Match`).
// Case (1b): We imported from the parent (`from typing.re import Match`,
// `Match`).
// Case (2): It's a builtin (like `list`).
if num_segments == 0
&& (module.is_empty()
|| from_imports
.get(module)
.map(|imports| imports.contains(member) || imports.contains("*"))
.unwrap_or(false))
{
return true; return true;
} }
if let Some((parent, member)) = target.rsplit_once('.') { // Case (3a): it's a fully qualified call path (`import typing`,
// Case (2): We imported star from the parent (`from typing.re import *`, // `typing.re.Match`). Case (3b): it's a fully qualified call path (`import
// `Match`). // typing.re`, `typing.re.Match`).
if call_path == member if num_segments > 0
&& from_imports && module
.get(parent) .split('.')
.map(|imports| imports.contains("*")) .enumerate()
.unwrap_or(false) .all(|(index, segment)| index < num_segments && call_path[index] == segment)
{ {
return true; return true;
}
// Case (3): We imported from the parent (`from typing.re import Match`,
// `Match`)
if call_path == member
&& from_imports
.get(parent)
.map(|imports| imports.contains(member))
.unwrap_or(false)
{
return true;
}
} }
// Case (4): We imported from the grandparent (`from typing import re`, // Case (4): We imported from the grandparent (`from typing import re`,
// `re.Match`) // `re.Match`)
let mut parts = target.rsplitn(3, '.'); if num_segments > 0 {
let member = parts.next(); // Find the number of common segments.
let parent = parts.next(); // TODO(charlie): Rewrite to avoid this allocation.
let grandparent = parts.next(); let parts: Vec<&str> = module.split('.').collect();
if let (Some(member), Some(parent), Some(grandparent)) = (member, parent, grandparent) { let num_matches = (0..parts.len())
if call_path == format!("{parent}.{member}") .take(num_segments)
&& from_imports .take_while(|i| parts[parts.len() - 1 - i] == call_path[num_segments - 1 - i])
.get(grandparent) .count();
.map(|imports| imports.contains(parent)) if num_matches > 0 {
// Verify that we have an import of the final common segment, from the remaining
// parent.
let cut = parts.len() - num_matches;
// TODO(charlie): Rewrite to avoid this allocation.
let module = parts[..cut].join(".");
let member = parts[cut];
if from_imports
.get(&module.as_str())
.map(|imports| imports.contains(member))
.unwrap_or(false) .unwrap_or(false)
{ {
return true; return true;
}
} }
} }
@ -197,12 +232,25 @@ mod tests {
use crate::ast::helpers::match_module_member; use crate::ast::helpers::match_module_member;
#[test]
fn builtin() -> Result<()> {
let expr = parser::parse_expression("list", "<filename>")?;
assert!(match_module_member(
&expr,
"",
"list",
&FnvHashMap::default()
));
Ok(())
}
#[test] #[test]
fn fully_qualified() -> Result<()> { fn fully_qualified() -> Result<()> {
let expr = parser::parse_expression("typing.re.Match", "<filename>")?; let expr = parser::parse_expression("typing.re.Match", "<filename>")?;
assert!(match_module_member( assert!(match_module_member(
&expr, &expr,
"typing.re.Match", "typing.re",
"Match",
&FnvHashMap::default() &FnvHashMap::default()
)); ));
Ok(()) Ok(())
@ -213,13 +261,15 @@ mod tests {
let expr = parser::parse_expression("Match", "<filename>")?; let expr = parser::parse_expression("Match", "<filename>")?;
assert!(!match_module_member( assert!(!match_module_member(
&expr, &expr,
"typing.re.Match", "typing.re",
"Match",
&FnvHashMap::default(), &FnvHashMap::default(),
)); ));
let expr = parser::parse_expression("re.Match", "<filename>")?; let expr = parser::parse_expression("re.Match", "<filename>")?;
assert!(!match_module_member( assert!(!match_module_member(
&expr, &expr,
"typing.re.Match", "typing.re",
"Match",
&FnvHashMap::default(), &FnvHashMap::default(),
)); ));
Ok(()) Ok(())
@ -230,7 +280,8 @@ mod tests {
let expr = parser::parse_expression("Match", "<filename>")?; let expr = parser::parse_expression("Match", "<filename>")?;
assert!(match_module_member( assert!(match_module_member(
&expr, &expr,
"typing.re.Match", "typing.re",
"Match",
&FnvHashMap::from_iter([("typing.re", FnvHashSet::from_iter(["*"]))]) &FnvHashMap::from_iter([("typing.re", FnvHashSet::from_iter(["*"]))])
)); ));
Ok(()) Ok(())
@ -241,7 +292,8 @@ mod tests {
let expr = parser::parse_expression("Match", "<filename>")?; let expr = parser::parse_expression("Match", "<filename>")?;
assert!(match_module_member( assert!(match_module_member(
&expr, &expr,
"typing.re.Match", "typing.re",
"Match",
&FnvHashMap::from_iter([("typing.re", FnvHashSet::from_iter(["Match"]))]) &FnvHashMap::from_iter([("typing.re", FnvHashSet::from_iter(["Match"]))])
)); ));
Ok(()) Ok(())
@ -252,7 +304,24 @@ mod tests {
let expr = parser::parse_expression("re.Match", "<filename>")?; let expr = parser::parse_expression("re.Match", "<filename>")?;
assert!(match_module_member( assert!(match_module_member(
&expr, &expr,
"typing.re.Match", "typing.re",
"Match",
&FnvHashMap::from_iter([("typing", FnvHashSet::from_iter(["re"]))])
));
let expr = parser::parse_expression("match.Match", "<filename>")?;
assert!(match_module_member(
&expr,
"typing.re.match",
"Match",
&FnvHashMap::from_iter([("typing.re", FnvHashSet::from_iter(["match"]))])
));
let expr = parser::parse_expression("re.match.Match", "<filename>")?;
assert!(match_module_member(
&expr,
"typing.re.match",
"Match",
&FnvHashMap::from_iter([("typing", FnvHashSet::from_iter(["re"]))]) &FnvHashMap::from_iter([("typing", FnvHashSet::from_iter(["re"]))])
)); ));
Ok(()) Ok(())

View file

@ -12,7 +12,7 @@ use rustpython_parser::ast::{
}; };
use rustpython_parser::parser; use rustpython_parser::parser;
use crate::ast::helpers::{extract_handler_names, match_module_member}; use crate::ast::helpers::{collect_call_paths, extract_handler_names, match_call_path};
use crate::ast::operations::extract_all_names; use crate::ast::operations::extract_all_names;
use crate::ast::relocate::relocate_expr; use crate::ast::relocate::relocate_expr;
use crate::ast::types::{ use crate::ast::types::{
@ -75,7 +75,7 @@ pub struct Checker<'a> {
seen_import_boundary: bool, seen_import_boundary: bool,
futures_allowed: bool, futures_allowed: bool,
annotations_future_enabled: bool, annotations_future_enabled: bool,
except_handlers: Vec<Vec<String>>, except_handlers: Vec<Vec<Vec<&'a str>>>,
} }
impl<'a> Checker<'a> { impl<'a> Checker<'a> {
@ -154,14 +154,10 @@ impl<'a> Checker<'a> {
} }
/// Return `true` if the `Expr` is a reference to `typing.${target}`. /// Return `true` if the `Expr` is a reference to `typing.${target}`.
pub fn match_typing_module(&self, expr: &Expr, target: &str) -> bool { pub fn match_typing_module(&self, call_path: &[&str], target: &str) -> bool {
match_module_member(expr, &format!("typing.{target}"), &self.from_imports) match_call_path(call_path, "typing", target, &self.from_imports)
|| (typing::in_extensions(target) || (typing::in_extensions(target)
&& match_module_member( && match_call_path(call_path, "typing_extensions", target, &self.from_imports))
expr,
&format!("typing_extensions.{target}"),
&self.from_imports,
))
} }
} }
@ -1046,7 +1042,7 @@ where
pyupgrade::plugins::use_pep604_annotation(self, expr, value, slice); pyupgrade::plugins::use_pep604_annotation(self, expr, value, slice);
} }
if self.match_typing_module(value, "Literal") { if self.match_typing_module(&collect_call_paths(value), "Literal") {
self.in_literal = true; self.in_literal = true;
} }
@ -1629,12 +1625,13 @@ where
args, args,
keywords, keywords,
} => { } => {
if self.match_typing_module(func, "ForwardRef") { let call_path = collect_call_paths(func);
if self.match_typing_module(&call_path, "ForwardRef") {
self.visit_expr(func); self.visit_expr(func);
for expr in args { for expr in args {
self.visit_annotation(expr); self.visit_annotation(expr);
} }
} else if self.match_typing_module(func, "cast") { } else if self.match_typing_module(&call_path, "cast") {
self.visit_expr(func); self.visit_expr(func);
if !args.is_empty() { if !args.is_empty() {
self.visit_annotation(&args[0]); self.visit_annotation(&args[0]);
@ -1642,12 +1639,12 @@ where
for expr in args.iter().skip(1) { for expr in args.iter().skip(1) {
self.visit_expr(expr); self.visit_expr(expr);
} }
} else if self.match_typing_module(func, "NewType") { } else if self.match_typing_module(&call_path, "NewType") {
self.visit_expr(func); self.visit_expr(func);
for expr in args.iter().skip(1) { for expr in args.iter().skip(1) {
self.visit_annotation(expr); self.visit_annotation(expr);
} }
} else if self.match_typing_module(func, "TypeVar") { } else if self.match_typing_module(&call_path, "TypeVar") {
self.visit_expr(func); self.visit_expr(func);
for expr in args.iter().skip(1) { for expr in args.iter().skip(1) {
self.visit_annotation(expr); self.visit_annotation(expr);
@ -1664,7 +1661,7 @@ where
} }
} }
} }
} else if self.match_typing_module(func, "NamedTuple") { } else if self.match_typing_module(&call_path, "NamedTuple") {
self.visit_expr(func); self.visit_expr(func);
// Ex) NamedTuple("a", [("a", int)]) // Ex) NamedTuple("a", [("a", int)])
@ -1696,7 +1693,7 @@ where
let KeywordData { value, .. } = &keyword.node; let KeywordData { value, .. } = &keyword.node;
self.visit_annotation(value); self.visit_annotation(value);
} }
} else if self.match_typing_module(func, "TypedDict") { } else if self.match_typing_module(&call_path, "TypedDict") {
self.visit_expr(func); self.visit_expr(func);
// Ex) TypedDict("a", {"a": int}) // Ex) TypedDict("a", {"a": int})
@ -2147,7 +2144,10 @@ impl<'a> Checker<'a> {
// Avoid flagging if NameError is handled. // Avoid flagging if NameError is handled.
if let Some(handler_names) = self.except_handlers.last() { if let Some(handler_names) = self.except_handlers.last() {
if handler_names.contains(&"NameError".to_string()) { if handler_names
.iter()
.any(|call_path| call_path.len() == 1 && call_path[0] == "NameError")
{
return; return;
} }
} }

View file

@ -7,7 +7,7 @@ use crate::check_ast::Checker;
use crate::checks::{Check, CheckCode, CheckKind}; use crate::checks::{Check, CheckCode, CheckKind};
fn is_sys(checker: &Checker, expr: &Expr, target: &str) -> bool { fn is_sys(checker: &Checker, expr: &Expr, target: &str) -> bool {
match_module_member(expr, &format!("sys.{target}"), &checker.from_imports) match_module_member(expr, "sys", target, &checker.from_imports)
} }
/// YTT101, YTT102, YTT301, YTT303 /// YTT101, YTT102, YTT301, YTT303
@ -181,7 +181,7 @@ pub fn compare(checker: &mut Checker, left: &Expr, ops: &[Cmpop], comparators: &
/// YTT202 /// YTT202
pub fn name_or_attribute(checker: &mut Checker, expr: &Expr) { pub fn name_or_attribute(checker: &mut Checker, expr: &Expr) {
if match_module_member(expr, "six.PY3", &checker.from_imports) { if match_module_member(expr, "six", "PY3", &checker.from_imports) {
checker.add_check(Check::new( checker.add_check(Check::new(
CheckKind::SixPY3Referenced, CheckKind::SixPY3Referenced,
Range::from_located(expr), Range::from_located(expr),

View file

@ -1,5 +1,6 @@
use rustpython_ast::{Arguments, Constant, Expr, ExprKind, Stmt, StmtKind}; use rustpython_ast::{Arguments, Constant, Expr, ExprKind, Stmt, StmtKind};
use crate::ast::helpers::collect_call_paths;
use crate::ast::types::Range; use crate::ast::types::Range;
use crate::ast::visitor; use crate::ast::visitor;
use crate::ast::visitor::Visitor; use crate::ast::visitor::Visitor;
@ -50,7 +51,7 @@ fn is_none_returning(body: &[Stmt]) -> bool {
/// ANN401 /// ANN401
fn check_dynamically_typed(checker: &mut Checker, annotation: &Expr, name: &str) { fn check_dynamically_typed(checker: &mut Checker, annotation: &Expr, name: &str) {
if checker.match_typing_module(annotation, "Any") { if checker.match_typing_module(&collect_call_paths(annotation), "Any") {
checker.add_check(Check::new( checker.add_check(Check::new(
CheckKind::DynamicallyTypedExpression(name.to_string()), CheckKind::DynamicallyTypedExpression(name.to_string()),
Range::from_located(annotation), Range::from_located(annotation),

View file

@ -1,7 +1,7 @@
use fnv::{FnvHashMap, FnvHashSet}; use fnv::{FnvHashMap, FnvHashSet};
use rustpython_ast::{Constant, Expr, ExprKind, Keyword, Stmt, StmtKind}; use rustpython_ast::{Constant, Expr, ExprKind, Keyword, Stmt, StmtKind};
use crate::ast::helpers::{compose_call_path, match_call_path}; use crate::ast::helpers::{collect_call_paths, match_call_path};
use crate::ast::types::Range; use crate::ast::types::Range;
use crate::check_ast::Checker; use crate::check_ast::Checker;
use crate::checks::{Check, CheckKind}; use crate::checks::{Check, CheckKind};
@ -18,14 +18,15 @@ fn is_abc_class(
.as_ref() .as_ref()
.map(|a| a == "metaclass") .map(|a| a == "metaclass")
.unwrap_or(false) .unwrap_or(false)
&& compose_call_path(&keyword.node.value) && match_call_path(
.map(|call_path| match_call_path(&call_path, "abc.ABCMeta", from_imports)) &collect_call_paths(&keyword.node.value),
.unwrap_or(false) "abc",
}) || bases.iter().any(|base| { "ABCMeta",
compose_call_path(base) from_imports,
.map(|call_path| match_call_path(&call_path, "abc.ABC", from_imports)) )
.unwrap_or(false) }) || bases
}) .iter()
.any(|base| match_call_path(&collect_call_paths(base), "abc", "ABC", from_imports))
} }
fn is_empty_body(body: &[Stmt]) -> bool { fn is_empty_body(body: &[Stmt]) -> bool {
@ -42,15 +43,21 @@ fn is_empty_body(body: &[Stmt]) -> bool {
} }
fn is_abstractmethod(expr: &Expr, from_imports: &FnvHashMap<&str, FnvHashSet<&str>>) -> bool { fn is_abstractmethod(expr: &Expr, from_imports: &FnvHashMap<&str, FnvHashSet<&str>>) -> bool {
compose_call_path(expr) match_call_path(
.map(|call_path| match_call_path(&call_path, "abc.abstractmethod", from_imports)) &collect_call_paths(expr),
.unwrap_or(false) "abc",
"abstractmethod",
from_imports,
)
} }
fn is_overload(expr: &Expr, from_imports: &FnvHashMap<&str, FnvHashSet<&str>>) -> bool { fn is_overload(expr: &Expr, from_imports: &FnvHashMap<&str, FnvHashSet<&str>>) -> bool {
compose_call_path(expr) match_call_path(
.map(|call_path| match_call_path(&call_path, "typing.overload", from_imports)) &collect_call_paths(expr),
.unwrap_or(false) "typing",
"overload",
from_imports,
)
} }
pub fn abstract_base_class( pub fn abstract_base_class(

View file

@ -10,10 +10,10 @@ pub fn assert_raises_exception(checker: &mut Checker, stmt: &Stmt, items: &[With
if let Some(item) = items.first() { if let Some(item) = items.first() {
let item_context = &item.context_expr; let item_context = &item.context_expr;
if let ExprKind::Call { func, args, .. } = &item_context.node { if let ExprKind::Call { func, args, .. } = &item_context.node {
if match_name_or_attr(func, "assertRaises") if args.len() == 1
&& args.len() == 1
&& match_name_or_attr(args.first().unwrap(), "Exception")
&& item.optional_vars.is_none() && item.optional_vars.is_none()
&& match_name_or_attr(func, "assertRaises")
&& match_name_or_attr(args.first().unwrap(), "Exception")
{ {
checker.add_check(Check::new( checker.add_check(Check::new(
CheckKind::NoAssertRaisesException, CheckKind::NoAssertRaisesException,

View file

@ -1,13 +1,14 @@
use rustpython_ast::{Expr, ExprKind}; use rustpython_ast::{Expr, ExprKind};
use crate::ast::helpers::{compose_call_path, match_module_member}; use crate::ast::helpers::{collect_call_paths, match_call_path};
use crate::ast::types::{Range, ScopeKind}; use crate::ast::types::{Range, ScopeKind};
use crate::check_ast::Checker; use crate::check_ast::Checker;
use crate::checks::{Check, CheckKind}; use crate::checks::{Check, CheckKind};
fn is_cache_func(checker: &Checker, expr: &Expr) -> bool { fn is_cache_func(checker: &Checker, expr: &Expr) -> bool {
match_module_member(expr, "functools.lru_cache", &checker.from_imports) let call_path = collect_call_paths(expr);
|| match_module_member(expr, "functools.cache", &checker.from_imports) match_call_path(&call_path, "functools", "lru_cache", &checker.from_imports)
|| match_call_path(&call_path, "functools", "cache", &checker.from_imports)
} }
/// B019 /// B019
@ -16,8 +17,8 @@ pub fn cached_instance_method(checker: &mut Checker, decorator_list: &[Expr]) {
for decorator in decorator_list { for decorator in decorator_list {
// TODO(charlie): This should take into account `classmethod-decorators` and // TODO(charlie): This should take into account `classmethod-decorators` and
// `staticmethod-decorators`. // `staticmethod-decorators`.
if let Some(decorator_path) = compose_call_path(decorator) { if let ExprKind::Name { id, .. } = &decorator.node {
if decorator_path == "classmethod" || decorator_path == "staticmethod" { if id == "classmethod" || id == "staticmethod" {
return; return;
} }
} }

View file

@ -1,7 +1,7 @@
use fnv::{FnvHashMap, FnvHashSet}; use fnv::{FnvHashMap, FnvHashSet};
use rustpython_ast::{Arguments, Constant, Expr, ExprKind}; use rustpython_ast::{Arguments, Constant, Expr, ExprKind};
use crate::ast::helpers::{compose_call_path, match_call_path}; use crate::ast::helpers::{collect_call_paths, compose_call_path, match_call_path};
use crate::ast::types::Range; use crate::ast::types::Range;
use crate::ast::visitor; use crate::ast::visitor;
use crate::ast::visitor::Visitor; use crate::ast::visitor::Visitor;
@ -9,34 +9,31 @@ use crate::check_ast::Checker;
use crate::checks::{Check, CheckKind}; use crate::checks::{Check, CheckKind};
use crate::flake8_bugbear::plugins::mutable_argument_default::is_mutable_func; use crate::flake8_bugbear::plugins::mutable_argument_default::is_mutable_func;
const IMMUTABLE_FUNCS: [&str; 7] = [ const IMMUTABLE_FUNCS: [(&str, &str); 7] = [
"tuple", ("", "tuple"),
"frozenset", ("", "frozenset"),
"operator.attrgetter", ("operator", "attrgetter"),
"operator.itemgetter", ("operator", "itemgetter"),
"operator.methodcaller", ("operator", "methodcaller"),
"types.MappingProxyType", ("types", "MappingProxyType"),
"re.compile", ("re", "compile"),
]; ];
fn is_immutable_func( fn is_immutable_func(
expr: &Expr, expr: &Expr,
extend_immutable_calls: &[&str], extend_immutable_calls: &[(&str, &str)],
from_imports: &FnvHashMap<&str, FnvHashSet<&str>>, from_imports: &FnvHashMap<&str, FnvHashSet<&str>>,
) -> bool { ) -> bool {
compose_call_path(expr) let call_path = collect_call_paths(expr);
.map(|call_path| { IMMUTABLE_FUNCS
IMMUTABLE_FUNCS .iter()
.iter() .chain(extend_immutable_calls)
.chain(extend_immutable_calls) .any(|(module, member)| match_call_path(&call_path, module, member, from_imports))
.any(|target| match_call_path(&call_path, target, from_imports))
})
.unwrap_or(false)
} }
struct ArgumentDefaultVisitor<'a> { struct ArgumentDefaultVisitor<'a> {
checks: Vec<(CheckKind, Range)>, checks: Vec<(CheckKind, Range)>,
extend_immutable_calls: &'a [&'a str], extend_immutable_calls: &'a [(&'a str, &'a str)],
from_imports: &'a FnvHashMap<&'a str, FnvHashSet<&'a str>>, from_imports: &'a FnvHashMap<&'a str, FnvHashSet<&'a str>>,
} }
@ -92,12 +89,20 @@ fn is_nan_or_infinity(expr: &Expr, args: &[Expr]) -> bool {
/// B008 /// B008
pub fn function_call_argument_default(checker: &mut Checker, arguments: &Arguments) { pub fn function_call_argument_default(checker: &mut Checker, arguments: &Arguments) {
let extend_immutable_cells: Vec<&str> = checker // Map immutable calls to (module, member) format.
let extend_immutable_cells: Vec<(&str, &str)> = checker
.settings .settings
.flake8_bugbear .flake8_bugbear
.extend_immutable_calls .extend_immutable_calls
.iter() .iter()
.map(|s| s.as_str()) .map(|s| {
let s = s.as_str();
if let Some(index) = s.rfind('.') {
(&s[..index], &s[index + 1..])
} else {
("", s)
}
})
.collect(); .collect();
let mut visitor = ArgumentDefaultVisitor { let mut visitor = ArgumentDefaultVisitor {
checks: vec![], checks: vec![],

View file

@ -1,29 +1,26 @@
use fnv::{FnvHashMap, FnvHashSet}; use fnv::{FnvHashMap, FnvHashSet};
use rustpython_ast::{Arguments, Expr, ExprKind}; use rustpython_ast::{Arguments, Expr, ExprKind};
use crate::ast::helpers::{compose_call_path, match_call_path}; use crate::ast::helpers::{collect_call_paths, match_call_path};
use crate::ast::types::Range; use crate::ast::types::Range;
use crate::check_ast::Checker; use crate::check_ast::Checker;
use crate::checks::{Check, CheckKind}; use crate::checks::{Check, CheckKind};
const MUTABLE_FUNCS: [&str; 7] = [ const MUTABLE_FUNCS: [(&str, &str); 7] = [
"dict", ("", "dict"),
"list", ("", "list"),
"set", ("", "set"),
"collections.Counter", ("collections", "Counter"),
"collections.OrderedDict", ("collections", "OrderedDict"),
"collections.defaultdict", ("collections", "defaultdict"),
"collections.deque", ("collections", "deque"),
]; ];
pub fn is_mutable_func(expr: &Expr, from_imports: &FnvHashMap<&str, FnvHashSet<&str>>) -> bool { pub fn is_mutable_func(expr: &Expr, from_imports: &FnvHashMap<&str, FnvHashSet<&str>>) -> bool {
compose_call_path(expr) let call_path = collect_call_paths(expr);
.map(|call_path| { MUTABLE_FUNCS
MUTABLE_FUNCS .iter()
.iter() .any(|(module, member)| match_call_path(&call_path, module, member, from_imports))
.any(|target| match_call_path(&call_path, target, from_imports))
})
.unwrap_or(false)
} }
/// B006 /// B006

View file

@ -1,16 +1,18 @@
use rustpython_ast::Expr; use rustpython_ast::Expr;
use crate::ast::helpers::{compose_call_path, match_call_path}; use crate::ast::helpers::{collect_call_paths, match_call_path};
use crate::ast::types::Range; use crate::ast::types::Range;
use crate::check_ast::Checker; use crate::check_ast::Checker;
use crate::checks::{Check, CheckKind}; use crate::checks::{Check, CheckKind};
/// B005 /// B005
pub fn useless_contextlib_suppress(checker: &mut Checker, expr: &Expr, args: &[Expr]) { pub fn useless_contextlib_suppress(checker: &mut Checker, expr: &Expr, args: &[Expr]) {
if compose_call_path(expr) if match_call_path(
.map(|call_path| match_call_path(&call_path, "contextlib.suppress", &checker.from_imports)) &collect_call_paths(expr),
.unwrap_or(false) "contextlib",
&& args.is_empty() "suppress",
&checker.from_imports,
) && args.is_empty()
{ {
checker.add_check(Check::new( checker.add_check(Check::new(
CheckKind::UselessContextlibSuppress, CheckKind::UselessContextlibSuppress,

View file

@ -2,7 +2,7 @@ use fnv::{FnvHashMap, FnvHashSet};
use itertools::Itertools; use itertools::Itertools;
use rustpython_ast::{Expr, ExprKind, Stmt, StmtKind}; use rustpython_ast::{Expr, ExprKind, Stmt, StmtKind};
use crate::ast::helpers::{compose_call_path, match_call_path, match_name_or_attr}; use crate::ast::helpers::{collect_call_paths, match_call_path, match_name_or_attr};
use crate::ast::types::{Scope, ScopeKind}; use crate::ast::types::{Scope, ScopeKind};
use crate::pep8_naming::settings::Settings; use crate::pep8_naming::settings::Settings;
use crate::python::string::{is_lower, is_upper}; use crate::python::string::{is_lower, is_upper};
@ -84,9 +84,12 @@ pub fn is_namedtuple_assignment(
from_imports: &FnvHashMap<&str, FnvHashSet<&str>>, from_imports: &FnvHashMap<&str, FnvHashSet<&str>>,
) -> bool { ) -> bool {
if let StmtKind::Assign { value, .. } = &stmt.node { if let StmtKind::Assign { value, .. } = &stmt.node {
compose_call_path(value) match_call_path(
.map(|call_path| match_call_path(&call_path, "collections.namedtuple", from_imports)) &collect_call_paths(value),
.unwrap_or(false) "collections",
"namedtuple",
from_imports,
)
} else { } else {
false false
} }

View file

@ -172,7 +172,8 @@ pub fn unnecessary_lru_cache_params(
keywords, keywords,
} = &expr.node } = &expr.node
{ {
if args.is_empty() && helpers::match_module_member(func, "functools.lru_cache", imports) if args.is_empty()
&& helpers::match_module_member(func, "functools", "lru_cache", imports)
{ {
// Ex) `functools.lru_cache()` // Ex) `functools.lru_cache()`
if keywords.is_empty() { if keywords.is_empty() {

View file

@ -1,5 +1,6 @@
use rustpython_ast::{Constant, Expr, ExprKind, Operator}; use rustpython_ast::{Constant, Expr, ExprKind, Operator};
use crate::ast::helpers::collect_call_paths;
use crate::ast::types::Range; use crate::ast::types::Range;
use crate::autofix::Fix; use crate::autofix::Fix;
use crate::check_ast::Checker; use crate::check_ast::Checker;
@ -43,7 +44,8 @@ fn union(elts: &[Expr]) -> Expr {
/// U007 /// U007
pub fn use_pep604_annotation(checker: &mut Checker, expr: &Expr, value: &Expr, slice: &Expr) { pub fn use_pep604_annotation(checker: &mut Checker, expr: &Expr, value: &Expr, slice: &Expr) {
if checker.match_typing_module(value, "Optional") { let call_path = collect_call_paths(value);
if checker.match_typing_module(&call_path, "Optional") {
let mut check = Check::new(CheckKind::UsePEP604Annotation, Range::from_located(expr)); let mut check = Check::new(CheckKind::UsePEP604Annotation, Range::from_located(expr));
if checker.patch() { if checker.patch() {
let mut generator = SourceGenerator::new(); let mut generator = SourceGenerator::new();
@ -58,7 +60,7 @@ pub fn use_pep604_annotation(checker: &mut Checker, expr: &Expr, value: &Expr, s
} }
} }
checker.add_check(check); checker.add_check(check);
} else if checker.match_typing_module(value, "Union") { } else if checker.match_typing_module(&call_path, "Union") {
let mut check = Check::new(CheckKind::UsePEP604Annotation, Range::from_located(expr)); let mut check = Check::new(CheckKind::UsePEP604Annotation, Range::from_located(expr));
if checker.patch() { if checker.patch() {
match &slice.node { match &slice.node {