From f67727b13c21b97360dbde75bf8007966cea4ce1 Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Mon, 14 Nov 2022 17:14:22 -0500 Subject: [PATCH] Improve performance of import matching code (#744) --- src/ast/helpers.rs | 185 ++++++++++++------ src/check_ast.rs | 34 ++-- src/flake8_2020/plugins.rs | 4 +- src/flake8_annotations/plugins.rs | 3 +- .../plugins/abstract_base_class.rs | 37 ++-- .../plugins/assert_raises_exception.rs | 6 +- .../plugins/cached_instance_method.rs | 11 +- .../plugins/function_call_argument_default.rs | 47 +++-- .../plugins/mutable_argument_default.rs | 29 ++- .../plugins/useless_contextlib_suppress.rs | 12 +- src/pep8_naming/helpers.rs | 11 +- src/pyupgrade/checks.rs | 3 +- .../plugins/use_pep604_annotation.rs | 6 +- 13 files changed, 238 insertions(+), 150 deletions(-) diff --git a/src/ast/helpers.rs b/src/ast/helpers.rs index ff6c91769d..df203e5c37 100644 --- a/src/ast/helpers.rs +++ b/src/ast/helpers.rs @@ -3,13 +3,14 @@ use once_cell::sync::Lazy; use regex::Regex; 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 { ExprKind::Call { func, .. } => { - compose_call_path_inner(func, parts); + collect_call_path_inner(func, parts); } ExprKind::Attribute { value, attr, .. } => { - compose_call_path_inner(value, parts); + collect_call_path_inner(value, parts); parts.push(attr); } 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`). +#[inline(always)] pub fn compose_call_path(expr: &Expr) -> Option { - let mut segments = vec![]; - compose_call_path_inner(expr, &mut segments); + let segments = collect_call_paths(expr); if segments.is_empty() { None } else { @@ -30,6 +31,14 @@ pub fn compose_call_path(expr: &Expr) -> Option { } } +/// 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}`. pub fn match_name_or_attr(expr: &Expr, target: &str) -> bool { match &expr.node { @@ -72,7 +81,7 @@ pub fn is_assignment_to_a_dunder(node: &StmtKind) -> bool { } /// Extract the names of all handled exceptions. -pub fn extract_handler_names(handlers: &[Excepthandler]) -> Vec { +pub fn extract_handler_names(handlers: &[Excepthandler]) -> Vec> { let mut handler_names = vec![]; for handler in handlers { match &handler.node { @@ -80,12 +89,16 @@ pub fn extract_handler_names(handlers: &[Excepthandler]) -> Vec { if let Some(type_) = type_ { if let ExprKind::Tuple { elts, .. } = &type_.node { for type_ in elts { - if let Some(name) = compose_call_path(type_) { - handler_names.push(name); + let call_path = collect_call_paths(type_); + if !call_path.is_empty() { + handler_names.push(call_path); } } - } else if let Some(name) = compose_call_path(type_) { - handler_names.push(name); + } else { + 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`. pub fn match_module_member( expr: &Expr, - target: &str, + module: &str, + member: &str, from_imports: &FnvHashMap<&str, FnvHashSet<&str>>, ) -> bool { - compose_call_path(expr) - .map(|expr| match_call_path(&expr, target, from_imports)) - .unwrap_or(false) + match_call_path(&collect_call_paths(expr), module, member, from_imports) } /// Return `true` if the `call_path` is a reference to `${module}.${target}`. /// /// Optimized version of `match_module_member` for pre-computed call paths. pub fn match_call_path( - call_path: &str, - target: &str, + call_path: &[&str], + module: &str, + member: &str, from_imports: &FnvHashMap<&str, FnvHashSet<&str>>, ) -> bool { - // Case (1a): it's the same call path (`import typing`, `typing.re.Match`). - // Case (1b): it's the same call path (`import typing.re`, `typing.re.Match`). - if call_path == target { + // If we have no segments, we can't ever match. + let num_segments = call_path.len(); + 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; } - if let Some((parent, member)) = target.rsplit_once('.') { - // Case (2): We imported star from the parent (`from typing.re import *`, - // `Match`). - if call_path == member - && from_imports - .get(parent) - .map(|imports| imports.contains("*")) - .unwrap_or(false) - { - 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 (3a): it's a fully qualified call path (`import typing`, + // `typing.re.Match`). Case (3b): it's a fully qualified call path (`import + // typing.re`, `typing.re.Match`). + if num_segments > 0 + && module + .split('.') + .enumerate() + .all(|(index, segment)| index < num_segments && call_path[index] == segment) + { + return true; } // Case (4): We imported from the grandparent (`from typing import re`, // `re.Match`) - let mut parts = target.rsplitn(3, '.'); - let member = parts.next(); - let parent = parts.next(); - let grandparent = parts.next(); - if let (Some(member), Some(parent), Some(grandparent)) = (member, parent, grandparent) { - if call_path == format!("{parent}.{member}") - && from_imports - .get(grandparent) - .map(|imports| imports.contains(parent)) + if num_segments > 0 { + // Find the number of common segments. + // TODO(charlie): Rewrite to avoid this allocation. + let parts: Vec<&str> = module.split('.').collect(); + let num_matches = (0..parts.len()) + .take(num_segments) + .take_while(|i| parts[parts.len() - 1 - i] == call_path[num_segments - 1 - i]) + .count(); + 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) - { - return true; + { + return true; + } } } @@ -197,12 +232,25 @@ mod tests { use crate::ast::helpers::match_module_member; + #[test] + fn builtin() -> Result<()> { + let expr = parser::parse_expression("list", "")?; + assert!(match_module_member( + &expr, + "", + "list", + &FnvHashMap::default() + )); + Ok(()) + } + #[test] fn fully_qualified() -> Result<()> { let expr = parser::parse_expression("typing.re.Match", "")?; assert!(match_module_member( &expr, - "typing.re.Match", + "typing.re", + "Match", &FnvHashMap::default() )); Ok(()) @@ -213,13 +261,15 @@ mod tests { let expr = parser::parse_expression("Match", "")?; assert!(!match_module_member( &expr, - "typing.re.Match", + "typing.re", + "Match", &FnvHashMap::default(), )); let expr = parser::parse_expression("re.Match", "")?; assert!(!match_module_member( &expr, - "typing.re.Match", + "typing.re", + "Match", &FnvHashMap::default(), )); Ok(()) @@ -230,7 +280,8 @@ mod tests { let expr = parser::parse_expression("Match", "")?; assert!(match_module_member( &expr, - "typing.re.Match", + "typing.re", + "Match", &FnvHashMap::from_iter([("typing.re", FnvHashSet::from_iter(["*"]))]) )); Ok(()) @@ -241,7 +292,8 @@ mod tests { let expr = parser::parse_expression("Match", "")?; assert!(match_module_member( &expr, - "typing.re.Match", + "typing.re", + "Match", &FnvHashMap::from_iter([("typing.re", FnvHashSet::from_iter(["Match"]))]) )); Ok(()) @@ -252,7 +304,24 @@ mod tests { let expr = parser::parse_expression("re.Match", "")?; assert!(match_module_member( &expr, - "typing.re.Match", + "typing.re", + "Match", + &FnvHashMap::from_iter([("typing", FnvHashSet::from_iter(["re"]))]) + )); + + let expr = parser::parse_expression("match.Match", "")?; + 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", "")?; + assert!(match_module_member( + &expr, + "typing.re.match", + "Match", &FnvHashMap::from_iter([("typing", FnvHashSet::from_iter(["re"]))]) )); Ok(()) diff --git a/src/check_ast.rs b/src/check_ast.rs index 6c15703601..e1cda22722 100644 --- a/src/check_ast.rs +++ b/src/check_ast.rs @@ -12,7 +12,7 @@ use rustpython_parser::ast::{ }; 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::relocate::relocate_expr; use crate::ast::types::{ @@ -75,7 +75,7 @@ pub struct Checker<'a> { seen_import_boundary: bool, futures_allowed: bool, annotations_future_enabled: bool, - except_handlers: Vec>, + except_handlers: Vec>>, } impl<'a> Checker<'a> { @@ -154,14 +154,10 @@ impl<'a> Checker<'a> { } /// Return `true` if the `Expr` is a reference to `typing.${target}`. - pub fn match_typing_module(&self, expr: &Expr, target: &str) -> bool { - match_module_member(expr, &format!("typing.{target}"), &self.from_imports) + pub fn match_typing_module(&self, call_path: &[&str], target: &str) -> bool { + match_call_path(call_path, "typing", target, &self.from_imports) || (typing::in_extensions(target) - && match_module_member( - expr, - &format!("typing_extensions.{target}"), - &self.from_imports, - )) + && match_call_path(call_path, "typing_extensions", target, &self.from_imports)) } } @@ -1046,7 +1042,7 @@ where 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; } @@ -1629,12 +1625,13 @@ where args, 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); for expr in args { 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); if !args.is_empty() { self.visit_annotation(&args[0]); @@ -1642,12 +1639,12 @@ where for expr in args.iter().skip(1) { 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); for expr in args.iter().skip(1) { 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); for expr in args.iter().skip(1) { 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); // Ex) NamedTuple("a", [("a", int)]) @@ -1696,7 +1693,7 @@ where let KeywordData { value, .. } = &keyword.node; 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); // Ex) TypedDict("a", {"a": int}) @@ -2147,7 +2144,10 @@ impl<'a> Checker<'a> { // Avoid flagging if NameError is handled. 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; } } diff --git a/src/flake8_2020/plugins.rs b/src/flake8_2020/plugins.rs index c5798d90f7..05feb813ff 100644 --- a/src/flake8_2020/plugins.rs +++ b/src/flake8_2020/plugins.rs @@ -7,7 +7,7 @@ use crate::check_ast::Checker; use crate::checks::{Check, CheckCode, CheckKind}; 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 @@ -181,7 +181,7 @@ pub fn compare(checker: &mut Checker, left: &Expr, ops: &[Cmpop], comparators: & /// YTT202 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( CheckKind::SixPY3Referenced, Range::from_located(expr), diff --git a/src/flake8_annotations/plugins.rs b/src/flake8_annotations/plugins.rs index 9d9914ac61..dc75ec8416 100644 --- a/src/flake8_annotations/plugins.rs +++ b/src/flake8_annotations/plugins.rs @@ -1,5 +1,6 @@ use rustpython_ast::{Arguments, Constant, Expr, ExprKind, Stmt, StmtKind}; +use crate::ast::helpers::collect_call_paths; use crate::ast::types::Range; use crate::ast::visitor; use crate::ast::visitor::Visitor; @@ -50,7 +51,7 @@ fn is_none_returning(body: &[Stmt]) -> bool { /// ANN401 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( CheckKind::DynamicallyTypedExpression(name.to_string()), Range::from_located(annotation), diff --git a/src/flake8_bugbear/plugins/abstract_base_class.rs b/src/flake8_bugbear/plugins/abstract_base_class.rs index eaa3e0e9ba..667b59cf38 100644 --- a/src/flake8_bugbear/plugins/abstract_base_class.rs +++ b/src/flake8_bugbear/plugins/abstract_base_class.rs @@ -1,7 +1,7 @@ use fnv::{FnvHashMap, FnvHashSet}; 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::check_ast::Checker; use crate::checks::{Check, CheckKind}; @@ -18,14 +18,15 @@ fn is_abc_class( .as_ref() .map(|a| a == "metaclass") .unwrap_or(false) - && compose_call_path(&keyword.node.value) - .map(|call_path| match_call_path(&call_path, "abc.ABCMeta", from_imports)) - .unwrap_or(false) - }) || bases.iter().any(|base| { - compose_call_path(base) - .map(|call_path| match_call_path(&call_path, "abc.ABC", from_imports)) - .unwrap_or(false) - }) + && match_call_path( + &collect_call_paths(&keyword.node.value), + "abc", + "ABCMeta", + from_imports, + ) + }) || bases + .iter() + .any(|base| match_call_path(&collect_call_paths(base), "abc", "ABC", from_imports)) } 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 { - compose_call_path(expr) - .map(|call_path| match_call_path(&call_path, "abc.abstractmethod", from_imports)) - .unwrap_or(false) + match_call_path( + &collect_call_paths(expr), + "abc", + "abstractmethod", + from_imports, + ) } fn is_overload(expr: &Expr, from_imports: &FnvHashMap<&str, FnvHashSet<&str>>) -> bool { - compose_call_path(expr) - .map(|call_path| match_call_path(&call_path, "typing.overload", from_imports)) - .unwrap_or(false) + match_call_path( + &collect_call_paths(expr), + "typing", + "overload", + from_imports, + ) } pub fn abstract_base_class( diff --git a/src/flake8_bugbear/plugins/assert_raises_exception.rs b/src/flake8_bugbear/plugins/assert_raises_exception.rs index c408f4b7b9..7fd17d8274 100644 --- a/src/flake8_bugbear/plugins/assert_raises_exception.rs +++ b/src/flake8_bugbear/plugins/assert_raises_exception.rs @@ -10,10 +10,10 @@ pub fn assert_raises_exception(checker: &mut Checker, stmt: &Stmt, items: &[With if let Some(item) = items.first() { let item_context = &item.context_expr; if let ExprKind::Call { func, args, .. } = &item_context.node { - if match_name_or_attr(func, "assertRaises") - && args.len() == 1 - && match_name_or_attr(args.first().unwrap(), "Exception") + if args.len() == 1 && item.optional_vars.is_none() + && match_name_or_attr(func, "assertRaises") + && match_name_or_attr(args.first().unwrap(), "Exception") { checker.add_check(Check::new( CheckKind::NoAssertRaisesException, diff --git a/src/flake8_bugbear/plugins/cached_instance_method.rs b/src/flake8_bugbear/plugins/cached_instance_method.rs index d205872525..298dd7f1bc 100644 --- a/src/flake8_bugbear/plugins/cached_instance_method.rs +++ b/src/flake8_bugbear/plugins/cached_instance_method.rs @@ -1,13 +1,14 @@ 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::check_ast::Checker; use crate::checks::{Check, CheckKind}; fn is_cache_func(checker: &Checker, expr: &Expr) -> bool { - match_module_member(expr, "functools.lru_cache", &checker.from_imports) - || match_module_member(expr, "functools.cache", &checker.from_imports) + let call_path = collect_call_paths(expr); + match_call_path(&call_path, "functools", "lru_cache", &checker.from_imports) + || match_call_path(&call_path, "functools", "cache", &checker.from_imports) } /// B019 @@ -16,8 +17,8 @@ pub fn cached_instance_method(checker: &mut Checker, decorator_list: &[Expr]) { for decorator in decorator_list { // TODO(charlie): This should take into account `classmethod-decorators` and // `staticmethod-decorators`. - if let Some(decorator_path) = compose_call_path(decorator) { - if decorator_path == "classmethod" || decorator_path == "staticmethod" { + if let ExprKind::Name { id, .. } = &decorator.node { + if id == "classmethod" || id == "staticmethod" { return; } } diff --git a/src/flake8_bugbear/plugins/function_call_argument_default.rs b/src/flake8_bugbear/plugins/function_call_argument_default.rs index 104f898aa1..3ae687badc 100644 --- a/src/flake8_bugbear/plugins/function_call_argument_default.rs +++ b/src/flake8_bugbear/plugins/function_call_argument_default.rs @@ -1,7 +1,7 @@ use fnv::{FnvHashMap, FnvHashSet}; 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::visitor; use crate::ast::visitor::Visitor; @@ -9,34 +9,31 @@ use crate::check_ast::Checker; use crate::checks::{Check, CheckKind}; use crate::flake8_bugbear::plugins::mutable_argument_default::is_mutable_func; -const IMMUTABLE_FUNCS: [&str; 7] = [ - "tuple", - "frozenset", - "operator.attrgetter", - "operator.itemgetter", - "operator.methodcaller", - "types.MappingProxyType", - "re.compile", +const IMMUTABLE_FUNCS: [(&str, &str); 7] = [ + ("", "tuple"), + ("", "frozenset"), + ("operator", "attrgetter"), + ("operator", "itemgetter"), + ("operator", "methodcaller"), + ("types", "MappingProxyType"), + ("re", "compile"), ]; fn is_immutable_func( expr: &Expr, - extend_immutable_calls: &[&str], + extend_immutable_calls: &[(&str, &str)], from_imports: &FnvHashMap<&str, FnvHashSet<&str>>, ) -> bool { - compose_call_path(expr) - .map(|call_path| { - IMMUTABLE_FUNCS - .iter() - .chain(extend_immutable_calls) - .any(|target| match_call_path(&call_path, target, from_imports)) - }) - .unwrap_or(false) + let call_path = collect_call_paths(expr); + IMMUTABLE_FUNCS + .iter() + .chain(extend_immutable_calls) + .any(|(module, member)| match_call_path(&call_path, module, member, from_imports)) } struct ArgumentDefaultVisitor<'a> { 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>>, } @@ -92,12 +89,20 @@ fn is_nan_or_infinity(expr: &Expr, args: &[Expr]) -> bool { /// B008 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 .flake8_bugbear .extend_immutable_calls .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(); let mut visitor = ArgumentDefaultVisitor { checks: vec![], diff --git a/src/flake8_bugbear/plugins/mutable_argument_default.rs b/src/flake8_bugbear/plugins/mutable_argument_default.rs index 94021f16ef..1ca00e4ba0 100644 --- a/src/flake8_bugbear/plugins/mutable_argument_default.rs +++ b/src/flake8_bugbear/plugins/mutable_argument_default.rs @@ -1,29 +1,26 @@ use fnv::{FnvHashMap, FnvHashSet}; 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::check_ast::Checker; use crate::checks::{Check, CheckKind}; -const MUTABLE_FUNCS: [&str; 7] = [ - "dict", - "list", - "set", - "collections.Counter", - "collections.OrderedDict", - "collections.defaultdict", - "collections.deque", +const MUTABLE_FUNCS: [(&str, &str); 7] = [ + ("", "dict"), + ("", "list"), + ("", "set"), + ("collections", "Counter"), + ("collections", "OrderedDict"), + ("collections", "defaultdict"), + ("collections", "deque"), ]; pub fn is_mutable_func(expr: &Expr, from_imports: &FnvHashMap<&str, FnvHashSet<&str>>) -> bool { - compose_call_path(expr) - .map(|call_path| { - MUTABLE_FUNCS - .iter() - .any(|target| match_call_path(&call_path, target, from_imports)) - }) - .unwrap_or(false) + let call_path = collect_call_paths(expr); + MUTABLE_FUNCS + .iter() + .any(|(module, member)| match_call_path(&call_path, module, member, from_imports)) } /// B006 diff --git a/src/flake8_bugbear/plugins/useless_contextlib_suppress.rs b/src/flake8_bugbear/plugins/useless_contextlib_suppress.rs index 7096901715..c5b3b37922 100644 --- a/src/flake8_bugbear/plugins/useless_contextlib_suppress.rs +++ b/src/flake8_bugbear/plugins/useless_contextlib_suppress.rs @@ -1,16 +1,18 @@ 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::check_ast::Checker; use crate::checks::{Check, CheckKind}; /// B005 pub fn useless_contextlib_suppress(checker: &mut Checker, expr: &Expr, args: &[Expr]) { - if compose_call_path(expr) - .map(|call_path| match_call_path(&call_path, "contextlib.suppress", &checker.from_imports)) - .unwrap_or(false) - && args.is_empty() + if match_call_path( + &collect_call_paths(expr), + "contextlib", + "suppress", + &checker.from_imports, + ) && args.is_empty() { checker.add_check(Check::new( CheckKind::UselessContextlibSuppress, diff --git a/src/pep8_naming/helpers.rs b/src/pep8_naming/helpers.rs index c819a66f43..04f3290cd0 100644 --- a/src/pep8_naming/helpers.rs +++ b/src/pep8_naming/helpers.rs @@ -2,7 +2,7 @@ use fnv::{FnvHashMap, FnvHashSet}; use itertools::Itertools; 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::pep8_naming::settings::Settings; use crate::python::string::{is_lower, is_upper}; @@ -84,9 +84,12 @@ pub fn is_namedtuple_assignment( from_imports: &FnvHashMap<&str, FnvHashSet<&str>>, ) -> bool { if let StmtKind::Assign { value, .. } = &stmt.node { - compose_call_path(value) - .map(|call_path| match_call_path(&call_path, "collections.namedtuple", from_imports)) - .unwrap_or(false) + match_call_path( + &collect_call_paths(value), + "collections", + "namedtuple", + from_imports, + ) } else { false } diff --git a/src/pyupgrade/checks.rs b/src/pyupgrade/checks.rs index cd08346353..9b1ee4d702 100644 --- a/src/pyupgrade/checks.rs +++ b/src/pyupgrade/checks.rs @@ -172,7 +172,8 @@ pub fn unnecessary_lru_cache_params( keywords, } = &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()` if keywords.is_empty() { diff --git a/src/pyupgrade/plugins/use_pep604_annotation.rs b/src/pyupgrade/plugins/use_pep604_annotation.rs index 9be264325f..762893c6e8 100644 --- a/src/pyupgrade/plugins/use_pep604_annotation.rs +++ b/src/pyupgrade/plugins/use_pep604_annotation.rs @@ -1,5 +1,6 @@ use rustpython_ast::{Constant, Expr, ExprKind, Operator}; +use crate::ast::helpers::collect_call_paths; use crate::ast::types::Range; use crate::autofix::Fix; use crate::check_ast::Checker; @@ -43,7 +44,8 @@ fn union(elts: &[Expr]) -> Expr { /// U007 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)); if checker.patch() { 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); - } 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)); if checker.patch() { match &slice.node {