Improve some import tracking code (#715)

2025-09-29 21:35:58 +00:00 · 2022-11-13 00:10:13 -05:00 · 2022-11-13 00:10:13 -05:00 · f5b1f957e3
commit f5b1f957e3
parent 8f99705795
10 changed files with 208 additions and 144 deletions
--- a/src/ast/helpers.rs
+++ b/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(())
+    }
+}
--- a/src/ast/operations.rs
+++ b/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 {
-    for index in parent_stack.iter().rev() {
-        let parent = parents[*index];
+pub fn on_conditional_branch<'a>(parents: &mut impl Iterator<Item = &'a Stmt>) -> bool {
+    parents.any(|parent| {
        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 {
-    for index in parent_stack.iter().rev() {
-        let parent = parents[*index];
-        if matches!(
+pub fn in_nested_block<'a>(parents: &mut impl Iterator<Item = &'a Stmt>) -> bool {
+    parents.any(|parent| {
+        matches!(
            parent.node,
            StmtKind::Try { .. } | StmtKind::If { .. } | StmtKind::With { .. }
-        ) {
-            return true;
-        }
-    }
-
-    false
+        )
+    })
 }

 /// Check if a node represents an unpacking assignment.
--- 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_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,
-                    "typing_extensions",
-                    self.from_imports.get("typing_extensions"),
+                    &format!("typing_extensions.{target}"),
+                    &self.from_imports,
                ))
    }
 }
@ -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;
            }

--- a/src/flake8_2020/plugins.rs
+++ b/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"))
-        && checker.settings.enabled.contains(&CheckCode::YTT202)
-    {
+    if match_module_member(expr, "six.PY3", &checker.from_imports) {
        checker.add_check(Check::new(
            CheckKind::SixPY3Referenced,
            Range::from_located(expr),
--- a/src/flake8_bugbear/plugins/cached_instance_method.rs
+++ b/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(
-        expr,
-        "lru_cache",
-        "functools",
-        checker.from_imports.get("functools"),
-    ) || match_name_or_attr_from_module(
-        expr,
-        "cache",
-        "functools",
-        checker.from_imports.get("functools"),
-    )
+    match_module_member(expr, "functools.lru_cache", &checker.from_imports)
+        || match_module_member(expr, "functools.cache", &checker.from_imports)
 }

 /// B019
--- 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;
+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(
-        || false,
-        |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))
+    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)
-                        {
-                            return true;
-                        }
-                    }
-                }
-            }
-
-            false
-        },
-    )
 }

 struct ArgumentDefaultVisitor<'a> {
--- a/src/flake8_bugbear/plugins/mutable_argument_default.rs
+++ b/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(
-        || 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))
+    compose_call_path(expr)
+        .map(|call_path| {
+            MUTABLE_FUNCS
+                .iter()
+                .any(|target| match_call_path(&call_path, target, from_imports))
+        })
        .unwrap_or(false)
-                        {
-                            return true;
-                        }
-                    }
-                }
-            }
-
-            false
-        },
-    )
 }

 /// B006
--- a/src/python/typing.rs
+++ b/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
-                .map(|imports| imports.contains(&id.as_str()))
+            from_imports
+                .get("typing")
+                .map(|imports| imports.contains(&id.as_str()) || imports.contains("*"))
                .unwrap_or_default()
                && PEP_585_BUILTINS_ELIGIBLE.contains(&id.as_str())
        }
--- a/src/pyupgrade/checks.rs
+++ b/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()
-                && helpers::match_name_or_attr_from_module(func, "lru_cache", "functools", imports)
+            if args.is_empty() && helpers::match_module_member(func, "functools.lru_cache", imports)
            {
                // Ex) `functools.lru_cache()`
                if keywords.is_empty() {
--- a/src/pyupgrade/plugins/unnecessary_lru_cache_params.rs
+++ b/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) =