Improve some import tracking code (#715)

src/ast/helpers.rs

@@ -1,4 +1,4 @@
-use fnv::FnvHashSet;
+use fnv::{FnvHashMap, FnvHashSet};
 use once_cell::sync::Lazy;
 use regex::Regex;
 use rustpython_ast::{Excepthandler, ExcepthandlerKind, Expr, ExprKind, Location, StmtKind};
@@ -19,6 +19,7 @@ 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`).
 pub fn compose_call_path(expr: &Expr) -> Option<String> {
     let mut segments = vec![];
     compose_call_path_inner(expr, &mut segments);
@@ -38,31 +39,6 @@ pub fn match_name_or_attr(expr: &Expr, target: &str) -> bool {
     }
 }
 
-/// Return `true` if the `Expr` is a reference to `${module}.${target}`.
-///
-/// Useful for, e.g., ensuring that a `Union` reference represents
-/// `typing.Union`.
-pub fn match_name_or_attr_from_module(
-    expr: &Expr,
-    target: &str,
-    module: &str,
-    imports: Option<&FnvHashSet<&str>>,
-) -> bool {
-    match &expr.node {
-        ExprKind::Attribute { value, attr, .. } => match &value.node {
-            ExprKind::Name { id, .. } => id == module && target == attr,
-            _ => false,
-        },
-        ExprKind::Name { id, .. } => {
-            target == id
-                && imports
-                    .map(|imports| imports.contains(&id.as_str()))
-                    .unwrap_or_default()
-        }
-        _ => false,
-    }
-}
-
 static DUNDER_REGEX: Lazy<Regex> = Lazy::new(|| Regex::new(r"__[^\s]+__").unwrap());
 
 pub fn is_assignment_to_a_dunder(node: &StmtKind) -> bool {
@@ -140,3 +116,145 @@ pub fn to_absolute(relative: &Location, base: &Location) -> Location {
         Location::new(relative.row() + base.row() - 1, relative.column())
     }
 }
+
+/// Return `true` if the `Expr` is a reference to `${module}.${target}`.
+///
+/// Useful for, e.g., ensuring that a `Union` reference represents
+/// `typing.Union`.
+pub fn match_module_member(
+    expr: &Expr,
+    target: &str,
+    from_imports: &FnvHashMap<&str, FnvHashSet<&str>>,
+) -> bool {
+    compose_call_path(expr)
+        .map(|expr| match_call_path(&expr, target, from_imports))
+        .unwrap_or(false)
+}
+
+/// 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,
+    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 {
+        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 (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))
+                .unwrap_or(false)
+        {
+            return true;
+        }
+    }
+
+    false
+}
+
+#[cfg(test)]
+mod tests {
+    use anyhow::Result;
+    use fnv::{FnvHashMap, FnvHashSet};
+    use rustpython_parser::parser;
+
+    use crate::ast::helpers::match_module_member;
+
+    #[test]
+    fn fully_qualified() -> Result<()> {
+        let expr = parser::parse_expression("typing.re.Match", "<filename>")?;
+        assert!(match_module_member(
+            &expr,
+            "typing.re.Match",
+            &FnvHashMap::default()
+        ));
+        Ok(())
+    }
+
+    #[test]
+    fn unimported() -> Result<()> {
+        let expr = parser::parse_expression("Match", "<filename>")?;
+        assert!(!match_module_member(
+            &expr,
+            "typing.re.Match",
+            &FnvHashMap::default(),
+        ));
+        let expr = parser::parse_expression("re.Match", "<filename>")?;
+        assert!(!match_module_member(
+            &expr,
+            "typing.re.Match",
+            &FnvHashMap::default(),
+        ));
+        Ok(())
+    }
+
+    #[test]
+    fn from_star() -> Result<()> {
+        let expr = parser::parse_expression("Match", "<filename>")?;
+        assert!(match_module_member(
+            &expr,
+            "typing.re.Match",
+            &FnvHashMap::from_iter([("typing.re", FnvHashSet::from_iter(["*"]))])
+        ));
+        Ok(())
+    }
+
+    #[test]
+    fn from_parent() -> Result<()> {
+        let expr = parser::parse_expression("Match", "<filename>")?;
+        assert!(match_module_member(
+            &expr,
+            "typing.re.Match",
+            &FnvHashMap::from_iter([("typing.re", FnvHashSet::from_iter(["Match"]))])
+        ));
+        Ok(())
+    }
+
+    #[test]
+    fn from_grandparent() -> Result<()> {
+        let expr = parser::parse_expression("re.Match", "<filename>")?;
+        assert!(match_module_member(
+            &expr,
+            "typing.re.Match",
+            &FnvHashMap::from_iter([("typing", FnvHashSet::from_iter(["re"]))])
+        ));
+        Ok(())
+    }
+}

src/ast/operations.rs

@@ -67,9 +67,8 @@ pub fn extract_all_names(stmt: &Stmt, scope: &Scope) -> Vec<String> {
 }
 
 /// Check if a node is parent of a conditional branch.
-pub fn on_conditional_branch(parent_stack: &[usize], parents: &[&Stmt]) -> bool {
+pub fn on_conditional_branch<'a>(parents: &mut impl Iterator<Item = &'a Stmt>) -> bool {
-    for index in parent_stack.iter().rev() {
+    parents.any(|parent| {
-        let parent = parents[*index];
         if matches!(parent.node, StmtKind::If { .. } | StmtKind::While { .. }) {
             return true;
         }
@@ -78,24 +77,18 @@ pub fn on_conditional_branch(parent_stack: &[usize], parents: &[&Stmt]) -> bool
                 return true;
             }
         }
-    }
-
         false
+    })
 }
 
 /// Check if a node is in a nested block.
-pub fn in_nested_block(parent_stack: &[usize], parents: &[&Stmt]) -> bool {
+pub fn in_nested_block<'a>(parents: &mut impl Iterator<Item = &'a Stmt>) -> bool {
-    for index in parent_stack.iter().rev() {
+    parents.any(|parent| {
-        let parent = parents[*index];
+        matches!(
-        if matches!(
             parent.node,
             StmtKind::Try { .. } | StmtKind::If { .. } | StmtKind::With { .. }
-        ) {
+        )
-            return true;
+    })
-        }
-    }
-
-    false
 }
 
 /// Check if a node represents an unpacking assignment.

src/check_ast.rs

@@ -12,7 +12,7 @@ use rustpython_parser::ast::{
 };
 use rustpython_parser::parser;
 
-use crate::ast::helpers::{extract_handler_names, match_name_or_attr_from_module};
+use crate::ast::helpers::{extract_handler_names, match_module_member};
 use crate::ast::operations::extract_all_names;
 use crate::ast::relocate::relocate_expr;
 use crate::ast::types::{
@@ -155,13 +155,12 @@ 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_name_or_attr_from_module(expr, target, "typing", self.from_imports.get("typing"))
+        match_module_member(expr, &format!("typing.{target}"), &self.from_imports)
             || (typing::in_extensions(target)
-                && match_name_or_attr_from_module(
+                && match_module_member(
                     expr,
-                    target,
+                    &format!("typing_extensions.{target}"),
-                    "typing_extensions",
+                    &self.from_imports,
-                    self.from_imports.get("typing_extensions"),
                 ))
     }
 }
@@ -192,7 +191,13 @@ where
                 self.futures_allowed = false;
                 if !self.seen_import_boundary
                     && !helpers::is_assignment_to_a_dunder(node)
-                    && !operations::in_nested_block(&self.parent_stack, &self.parents)
+                    && !operations::in_nested_block(
+                        &mut self
+                            .parent_stack
+                            .iter()
+                            .rev()
+                            .map(|index| self.parents[*index]),
+                    )
                 {
                     self.seen_import_boundary = true;
                 }
@@ -1019,7 +1024,7 @@ where
                         // Ex) List[...]
                         if self.settings.enabled.contains(&CheckCode::U006)
                             && self.settings.target_version >= PythonVersion::Py39
-                            && typing::is_pep585_builtin(expr, self.from_imports.get("typing"))
+                            && typing::is_pep585_builtin(expr, &self.from_imports)
                         {
                             pyupgrade::plugins::use_pep585_annotation(self, expr, id);
                         }
@@ -1051,7 +1056,7 @@ where
                 // Ex) typing.List[...]
                 if self.settings.enabled.contains(&CheckCode::U006)
                     && self.settings.target_version >= PythonVersion::Py39
-                    && typing::is_pep585_builtin(expr, self.from_imports.get("typing"))
+                    && typing::is_pep585_builtin(expr, &self.from_imports)
                 {
                     pyupgrade::plugins::use_pep585_annotation(self, expr, attr);
                 }
@@ -2214,7 +2219,13 @@ impl<'a> Checker<'a> {
 
     fn handle_node_delete(&mut self, expr: &Expr) {
         if let ExprKind::Name { id, .. } = &expr.node {
-            if operations::on_conditional_branch(&self.parent_stack, &self.parents) {
+            if operations::on_conditional_branch(
+                &mut self
+                    .parent_stack
+                    .iter()
+                    .rev()
+                    .map(|index| self.parents[*index]),
+            ) {
                 return;
             }
 

src/flake8_2020/plugins.rs

@@ -1,13 +1,13 @@
 use num_bigint::BigInt;
 use rustpython_ast::{Cmpop, Constant, Expr, ExprKind, Located};
 
-use crate::ast::helpers::match_name_or_attr_from_module;
+use crate::ast::helpers::match_module_member;
 use crate::ast::types::Range;
 use crate::check_ast::Checker;
 use crate::checks::{Check, CheckCode, CheckKind};
 
 fn is_sys(checker: &Checker, expr: &Expr, target: &str) -> bool {
-    match_name_or_attr_from_module(expr, target, "sys", checker.from_imports.get("sys"))
+    match_module_member(expr, &format!("sys.{target}"), &checker.from_imports)
 }
 
 /// YTT101, YTT102, YTT301, YTT303
@@ -181,9 +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_name_or_attr_from_module(expr, "PY3", "six", checker.from_imports.get("six"))
+    if match_module_member(expr, "six.PY3", &checker.from_imports) {
-        && checker.settings.enabled.contains(&CheckCode::YTT202)
-    {
         checker.add_check(Check::new(
             CheckKind::SixPY3Referenced,
             Range::from_located(expr),

src/flake8_bugbear/plugins/cached_instance_method.rs

@@ -1,22 +1,13 @@
 use rustpython_ast::{Expr, ExprKind};
 
-use crate::ast::helpers::{compose_call_path, match_name_or_attr_from_module};
+use crate::ast::helpers::{compose_call_path, match_module_member};
 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_name_or_attr_from_module(
+    match_module_member(expr, "functools.lru_cache", &checker.from_imports)
-        expr,
+        || match_module_member(expr, "functools.cache", &checker.from_imports)
-        "lru_cache",
-        "functools",
-        checker.from_imports.get("functools"),
-    ) || match_name_or_attr_from_module(
-        expr,
-        "cache",
-        "functools",
-        checker.from_imports.get("functools"),
-    )
 }
 
 /// B019

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;
+use crate::ast::helpers::{compose_call_path, match_call_path};
 use crate::ast::types::Range;
 use crate::ast::visitor;
 use crate::ast::visitor::Visitor;
@@ -24,37 +24,14 @@ fn is_immutable_func(
     extend_immutable_calls: &[&str],
     from_imports: &FnvHashMap<&str, FnvHashSet<&str>>,
 ) -> bool {
-    compose_call_path(expr).map_or_else(
+    compose_call_path(expr)
-        || false,
+        .map(|call_path| {
-        |call_path| {
+            IMMUTABLE_FUNCS
-            // It matches the call path exactly (`operator.methodcaller`).
+                .iter()
-            for target in IMMUTABLE_FUNCS.iter().chain(extend_immutable_calls) {
+                .chain(extend_immutable_calls)
-                if &call_path == target {
+                .any(|target| match_call_path(&call_path, target, from_imports))
-                    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> {

src/flake8_bugbear/plugins/mutable_argument_default.rs

@@ -1,7 +1,7 @@
 use fnv::{FnvHashMap, FnvHashSet};
 use rustpython_ast::{Arguments, Expr, ExprKind};
 
-use crate::ast::helpers::compose_call_path;
+use crate::ast::helpers::{compose_call_path, match_call_path};
 use crate::ast::types::Range;
 use crate::check_ast::Checker;
 use crate::checks::{Check, CheckKind};
@@ -17,37 +17,13 @@ const MUTABLE_FUNCS: [&str; 7] = [
 ];
 
 pub fn is_mutable_func(expr: &Expr, from_imports: &FnvHashMap<&str, FnvHashSet<&str>>) -> bool {
-    compose_call_path(expr).map_or_else(
+    compose_call_path(expr)
-        || false,
+        .map(|call_path| {
-        |call_path| {
+            MUTABLE_FUNCS
-            // It matches the call path exactly (`collections.Counter`).
+                .iter()
-            for target in MUTABLE_FUNCS {
+                .any(|target| match_call_path(&call_path, target, from_imports))
-                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)
-                        {
-                            return true;
-                        }
-                    }
-                }
-            }
-
-            false
-        },
-    )
 }
 
 /// B006

src/python/typing.rs

@@ -209,7 +209,7 @@ pub enum SubscriptKind {
 
 pub fn match_annotated_subscript(
     expr: &Expr,
-    imports: &FnvHashMap<&str, FnvHashSet<&str>>,
+    from_imports: &FnvHashMap<&str, FnvHashSet<&str>>,
 ) -> Option<SubscriptKind> {
     match &expr.node {
         ExprKind::Attribute { attr, value, .. } => {
@@ -238,9 +238,9 @@ pub fn match_annotated_subscript(
             // Verify that, e.g., `Union` is a reference to `typing.Union`.
             if let Some(modules) = IMPORTED_SUBSCRIPTS.get(&id.as_str()) {
                 for module in modules {
-                    if imports
+                    if from_imports
                         .get(module)
-                        .map(|imports| imports.contains(&id.as_str()))
+                        .map(|imports| imports.contains(&id.as_str()) || imports.contains("*"))
                         .unwrap_or_default()
                     {
                         return if is_pep593_annotated_subscript(id) {
@@ -260,7 +260,7 @@ pub fn match_annotated_subscript(
 /// Returns `true` if `Expr` represents a reference to a typing object with a
 /// PEP 585 built-in. Note that none of the PEP 585 built-ins are in
 /// `typing_extensions`.
-pub fn is_pep585_builtin(expr: &Expr, typing_imports: Option<&FnvHashSet<&str>>) -> bool {
+pub fn is_pep585_builtin(expr: &Expr, from_imports: &FnvHashMap<&str, FnvHashSet<&str>>) -> bool {
     match &expr.node {
         ExprKind::Attribute { attr, value, .. } => {
             if let ExprKind::Name { id, .. } = &value.node {
@@ -270,8 +270,9 @@ pub fn is_pep585_builtin(expr: &Expr, typing_imports: Option<&FnvHashSet<&str>>)
             }
         }
         ExprKind::Name { id, .. } => {
-            typing_imports
+            from_imports
-                .map(|imports| imports.contains(&id.as_str()))
+                .get("typing")
+                .map(|imports| imports.contains(&id.as_str()) || imports.contains("*"))
                 .unwrap_or_default()
                 && PEP_585_BUILTINS_ELIGIBLE.contains(&id.as_str())
         }

src/pyupgrade/checks.rs

@@ -1,4 +1,4 @@
-use fnv::FnvHashSet;
+use fnv::{FnvHashMap, FnvHashSet};
 use rustpython_ast::{Constant, KeywordData};
 use rustpython_parser::ast::{ArgData, Expr, ExprKind, Stmt, StmtKind};
 
@@ -163,7 +163,7 @@ pub fn type_of_primitive(func: &Expr, args: &[Expr], location: Range) -> Option<
 pub fn unnecessary_lru_cache_params(
     decorator_list: &[Expr],
     target_version: PythonVersion,
-    imports: Option<&FnvHashSet<&str>>,
+    imports: &FnvHashMap<&str, FnvHashSet<&str>>,
 ) -> Option<Check> {
     for expr in decorator_list.iter() {
         if let ExprKind::Call {
@@ -172,8 +172,7 @@ pub fn unnecessary_lru_cache_params(
             keywords,
         } = &expr.node
         {
-            if args.is_empty()
+            if args.is_empty() && helpers::match_module_member(func, "functools.lru_cache", imports)
-                && helpers::match_name_or_attr_from_module(func, "lru_cache", "functools", imports)
             {
                 // Ex) `functools.lru_cache()`
                 if keywords.is_empty() {

src/pyupgrade/plugins/unnecessary_lru_cache_params.rs

@@ -8,7 +8,7 @@ pub fn unnecessary_lru_cache_params(checker: &mut Checker, decorator_list: &[Exp
     if let Some(mut check) = checks::unnecessary_lru_cache_params(
         decorator_list,
         checker.settings.target_version,
-        checker.from_imports.get("functools"),
+        &checker.from_imports,
     ) {
         if checker.patch() {
             if let Some(fix) =