Validate that mutable and immutable defaults are imported (#710)

resources/test/fixtures/B008_extended.py

@@ -4,14 +4,17 @@ import fastapi
 from fastapi import Query
 
 
-def this_is_okay_extended(db=fastapi.Depends(get_db)):
+def okay(db=fastapi.Depends(get_db)):
     ...
 
 
-def this_is_okay_extended_second(data: List[str] = fastapi.Query(None)):
+def okay(data: List[str] = fastapi.Query(None)):
     ...
 
 
-# TODO(charlie): Support `import from`.
-def this_is_not_okay_relative_import_not_listed(data: List[str] = Query(None)):
+def okay(data: List[str] = Query(None)):
+    ...
+
+
+def error_due_to_missing_import(data: List[str] = Depends(None)):
     ...

src/flake8_bugbear/plugins/function_call_argument_default.rs

@@ -1,3 +1,4 @@
+use fnv::{FnvHashMap, FnvHashSet};
 use rustpython_ast::{Arguments, Constant, Expr, ExprKind};
 
 use crate::ast::helpers::compose_call_path;
@@ -8,31 +9,58 @@ use crate::check_ast::Checker;
 use crate::checks::{Check, CheckKind};
 use crate::flake8_bugbear::plugins::mutable_argument_default::is_mutable_func;
 
-// TODO(charlie): Verify imports for each of the imported members.
-const IMMUTABLE_FUNCS: [&str; 11] = [
+const IMMUTABLE_FUNCS: [&str; 7] = [
     "tuple",
     "frozenset",
     "operator.attrgetter",
     "operator.itemgetter",
     "operator.methodcaller",
-    "attrgetter",
-    "itemgetter",
-    "methodcaller",
     "types.MappingProxyType",
-    "MappingProxyType",
     "re.compile",
 ];
 
-fn is_immutable_func(expr: &Expr, extend_immutable_calls: &[String]) -> bool {
+fn is_immutable_func(
+    expr: &Expr,
+    extend_immutable_calls: &[&str],
+    from_imports: &FnvHashMap<&str, FnvHashSet<&str>>,
+) -> bool {
     compose_call_path(expr).map_or_else(
         || false,
-        |func| IMMUTABLE_FUNCS.contains(&func.as_str()) || extend_immutable_calls.contains(&func),
+        |call_path| {
+            // It matches the call path exactly (`operator.methodcaller`).
+            for target in IMMUTABLE_FUNCS.iter().chain(extend_immutable_calls) {
+                if &call_path == target {
+                    return true;
+                }
+            }
+
+            // It matches the member name, and was imported from that module (`methodcaller`
+            // following `from operator import methodcaller`).
+            if !call_path.contains('.') {
+                for target in IMMUTABLE_FUNCS.iter().chain(extend_immutable_calls) {
+                    let mut splitter = target.rsplit('.');
+                    if let (Some(member), Some(module)) = (splitter.next(), splitter.next()) {
+                        if call_path == member
+                            && from_imports
+                                .get(module)
+                                .map(|module| module.contains(member))
+                                .unwrap_or(false)
+                        {
+                            return true;
+                        }
+                    }
+                }
+            }
+
+            false
+        },
     )
 }
 
 struct ArgumentDefaultVisitor<'a> {
     checks: Vec<(CheckKind, Range)>,
-    extend_immutable_calls: &'a [String],
+    extend_immutable_calls: &'a [&'a str],
+    from_imports: &'a FnvHashMap<&'a str, FnvHashSet<&'a str>>,
 }
 
 impl<'a, 'b> Visitor<'b> for ArgumentDefaultVisitor<'b>
@@ -42,8 +70,8 @@ where
     fn visit_expr(&mut self, expr: &'b Expr) {
         match &expr.node {
             ExprKind::Call { func, args, .. } => {
-                if !is_mutable_func(func)
-                    && !is_immutable_func(func, self.extend_immutable_calls)
+                if !is_mutable_func(func, self.from_imports)
+                    && !is_immutable_func(func, self.extend_immutable_calls, self.from_imports)
                     && !is_nan_or_infinity(func, args)
                 {
                     self.checks.push((
@@ -87,9 +115,17 @@ 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
+        .settings
+        .flake8_bugbear
+        .extend_immutable_calls
+        .iter()
+        .map(|s| s.as_str())
+        .collect();
     let mut visitor = ArgumentDefaultVisitor {
         checks: vec![],
-        extend_immutable_calls: &checker.settings.flake8_bugbear.extend_immutable_calls,
+        extend_immutable_calls: &extend_immutable_cells,
+        from_imports: &checker.from_imports,
     };
     for expr in arguments
         .defaults

src/flake8_bugbear/plugins/mutable_argument_default.rs

@@ -1,33 +1,53 @@
+use crate::ast::helpers::compose_call_path;
+use fnv::{FnvHashMap, FnvHashSet};
 use rustpython_ast::{Arguments, Expr, ExprKind};
 
 use crate::ast::types::Range;
 use crate::check_ast::Checker;
 use crate::checks::{Check, CheckKind};
 
-// TODO(charlie): Verify imports for each of the imported members.
-pub fn is_mutable_func(expr: &Expr) -> bool {
-    match &expr.node {
-        ExprKind::Name { id, .. }
-            if id == "dict"
-                || id == "list"
-                || id == "set"
-                || id == "Counter"
-                || id == "OrderedDict"
-                || id == "defaultdict"
-                || id == "deque" =>
+const MUTABLE_FUNCS: [&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_or_else(
+        || false,
+        |call_path| {
+            // It matches the call path exactly (`collections.Counter`).
+            for target in MUTABLE_FUNCS {
+                if call_path == target {
+                    return true;
+                }
+            }
+
+            // It matches the member name, and was imported from that module (`Counter`
+            // following `from collections import Counter`).
+            if !call_path.contains('.') {
+                for target in MUTABLE_FUNCS {
+                    let mut splitter = target.rsplit('.');
+                    if let (Some(member), Some(module)) = (splitter.next(), splitter.next()) {
+                        if call_path == member
+                            && from_imports
+                                .get(module)
+                                .map(|module| module.contains(member))
+                                .unwrap_or(false)
                         {
-            true
+                            return true;
                         }
-        ExprKind::Attribute { value, attr, .. }
-            if (attr == "Counter"
-                || attr == "OrderedDict"
-                || attr == "defaultdict"
-                || attr == "deque") =>
-        {
-            matches!(&value.node, ExprKind::Name { id, .. } if id == "collections")
                     }
-        _ => false,
                 }
+            }
+
+            false
+        },
+    )
 }
 
 /// B006
@@ -50,7 +70,7 @@ pub fn mutable_argument_default(checker: &mut Checker, arguments: &Arguments) {
                 ));
             }
             ExprKind::Call { func, .. } => {
-                if is_mutable_func(func) {
+                if is_mutable_func(func, &checker.from_imports) {
                     checker.add_check(Check::new(
                         CheckKind::MutableArgumentDefault,
                         Range::from_located(expr),

src/flake8_bugbear/snapshots/ruff__flake8_bugbear__tests__extend_immutable_calls.snap

@@ -4,10 +4,10 @@ expression: checks
 ---
 - kind: FunctionCallArgumentDefault
   location:
-    row: 16
-    column: 66
+    row: 19
+    column: 50
   end_location:
-    row: 16
-    column: 77
+    row: 19
+    column: 63
   fix: ~