Add import name resolution to Context (#3777)

crates/ruff/resources/test/fixtures/pylint/sys_exit_alias_3.py

@@ -7,3 +7,10 @@ quit(0)
 def main():
     exit(1)
     quit(1)
+
+
+def main():
+    exit = 1
+
+    exit(1)
+    quit(1)

crates/ruff/resources/test/fixtures/pyupgrade/UP033_0.py

@@ -1,5 +1,4 @@
 import functools
-from functools import lru_cache
 
 
 @functools.lru_cache(maxsize=None)
@@ -7,11 +6,6 @@ def fixme():
     pass
 
 
-@lru_cache(maxsize=None)
-def fixme():
-    pass
-
-
 @other_decorator
 @functools.lru_cache(maxsize=None)
 def fixme():
@@ -29,31 +23,16 @@ def ok():
     pass
 
 
-@lru_cache()
-def ok():
-    pass
-
-
 @functools.lru_cache(maxsize=64)
 def ok():
     pass
 
 
-@lru_cache(maxsize=64)
-def ok():
-    pass
-
-
 def user_func():
     pass
 
 
-@lru_cache(user_func)
+@functools.lru_cache(user_func)
-def ok():
-    pass
-
-
-@lru_cache(user_func, maxsize=None)
 def ok():
     pass
 

crates/ruff/resources/test/fixtures/pyupgrade/UP033_1.py

@@ -0,0 +1,51 @@
+from functools import lru_cache
+
+
+@lru_cache(maxsize=None)
+def fixme():
+    pass
+
+
+@other_decorator
+@lru_cache(maxsize=None)
+def fixme():
+    pass
+
+
+@lru_cache(maxsize=None)
+@other_decorator
+def fixme():
+    pass
+
+
+@lru_cache()
+def ok():
+    pass
+
+
+@lru_cache(maxsize=64)
+def ok():
+    pass
+
+
+def user_func():
+    pass
+
+
+@lru_cache(user_func)
+def ok():
+    pass
+
+
+@lru_cache(user_func, maxsize=None)
+def ok():
+    pass
+
+
+def lru_cache(maxsize=None):
+    pass
+
+
+@lru_cache(maxsize=None)
+def ok():
+    pass

crates/ruff/src/rules/pylint/rules/sys_exit_alias.rs

@@ -2,8 +2,6 @@ use rustpython_parser::ast::{Expr, ExprKind};
 
 use ruff_diagnostics::{AutofixKind, Diagnostic, Edit, Violation};
 use ruff_macros::{derive_message_formats, violation};
-use ruff_python_ast::context::Context;
-use ruff_python_ast::scope::{BindingKind, FromImportation, Importation, SubmoduleImportation};
 use ruff_python_ast::types::Range;
 
 use crate::checkers::ast::Checker;
@@ -28,52 +26,6 @@ impl Violation for SysExitAlias {
     }
 }
 
-/// Return the appropriate `sys.exit` reference based on the current set of
-/// imports, or `None` is `sys.exit` hasn't been imported.
-fn get_member_import_name_alias(context: &Context, module: &str, member: &str) -> Option<String> {
-    context.scopes().find_map(|scope| {
-        scope
-            .binding_ids()
-            .find_map(|index| match &context.bindings[*index].kind {
-                // e.g. module=sys object=exit
-                // `import sys`         -> `sys.exit`
-                // `import sys as sys2` -> `sys2.exit`
-                BindingKind::Importation(Importation { name, full_name }) => {
-                    if full_name == &module {
-                        Some(format!("{name}.{member}"))
-                    } else {
-                        None
-                    }
-                }
-                // e.g. module=os.path object=join
-                // `from os.path import join`          -> `join`
-                // `from os.path import join as join2` -> `join2`
-                BindingKind::FromImportation(FromImportation { name, full_name }) => {
-                    let mut parts = full_name.split('.');
-                    if parts.next() == Some(module)
-                        && parts.next() == Some(member)
-                        && parts.next().is_none()
-                    {
-                        Some((*name).to_string())
-                    } else {
-                        None
-                    }
-                }
-                // e.g. module=os.path object=join
-                // `import os.path ` -> `os.path.join`
-                BindingKind::SubmoduleImportation(SubmoduleImportation { full_name, .. }) => {
-                    if full_name == &module {
-                        Some(format!("{full_name}.{member}"))
-                    } else {
-                        None
-                    }
-                }
-                // Non-imports.
-                _ => None,
-            })
-    })
-}
-
 /// PLR1722
 pub fn sys_exit_alias(checker: &mut Checker, func: &Expr) {
     let ExprKind::Name { id, .. } = &func.node else {
@@ -93,9 +45,9 @@ pub fn sys_exit_alias(checker: &mut Checker, func: &Expr) {
             Range::from(func),
         );
         if checker.patch(diagnostic.kind.rule()) {
-            if let Some(content) = get_member_import_name_alias(&checker.ctx, "sys", "exit") {
+            if let Some(binding) = checker.ctx.resolve_qualified_import_name("sys", "exit") {
                 diagnostic.set_fix(Edit::replacement(
-                    content,
+                    binding,
                     func.location,
                     func.end_location.unwrap(),
                 ));

crates/ruff/src/rules/pylint/snapshots/ruff__rules__pylint__tests__PLR1722_sys_exit_alias_3.py.snap

@@ -44,4 +44,18 @@ expression: diagnostics
           row: 9
           column: 8
   parent: ~
+- kind:
+    name: SysExitAlias
+    body: "Use `sys.exit()` instead of `quit`"
+    suggestion: "Replace `quit` with `sys.exit()`"
+    fixable: true
+  location:
+    row: 16
+    column: 4
+  end_location:
+    row: 16
+    column: 8
+  fix:
+    edits: []
+  parent: ~
 

crates/ruff/src/rules/pyupgrade/mod.rs

@@ -59,7 +59,8 @@ mod tests {
     #[test_case(Rule::PrintfStringFormatting, Path::new("UP031_0.py"); "UP031_0")]
     #[test_case(Rule::PrintfStringFormatting, Path::new("UP031_1.py"); "UP031_1")]
     #[test_case(Rule::FString, Path::new("UP032.py"); "UP032")]
-    #[test_case(Rule::LRUCacheWithMaxsizeNone, Path::new("UP033.py"); "UP033")]
+    #[test_case(Rule::LRUCacheWithMaxsizeNone, Path::new("UP033_0.py"); "UP033_0")]
+    #[test_case(Rule::LRUCacheWithMaxsizeNone, Path::new("UP033_1.py"); "UP033_1")]
     #[test_case(Rule::ExtraneousParentheses, Path::new("UP034.py"); "UP034")]
     #[test_case(Rule::DeprecatedImport, Path::new("UP035.py"); "UP035")]
     #[test_case(Rule::OutdatedVersionBlock, Path::new("UP036_0.py"); "UP036_0")]

crates/ruff/src/rules/pyupgrade/rules/lru_cache_with_maxsize_none.rs

@@ -2,7 +2,6 @@ use rustpython_parser::ast::{Constant, Expr, ExprKind, KeywordData};
 
 use ruff_diagnostics::{AlwaysAutofixableViolation, Diagnostic, Edit};
 use ruff_macros::{derive_message_formats, violation};
-use ruff_python_ast::helpers::{create_expr, unparse_expr};
 use ruff_python_ast::types::Range;
 
 use crate::checkers::ast::Checker;
@@ -58,16 +57,12 @@ pub fn lru_cache_with_maxsize_none(checker: &mut Checker, decorator_list: &[Expr
                     Range::new(func.end_location.unwrap(), expr.end_location.unwrap()),
                 );
                 if checker.patch(diagnostic.kind.rule()) {
-                    if let ExprKind::Attribute { value, ctx, .. } = &func.node {
+                    if let Some(binding) = checker
+                        .ctx
+                        .resolve_qualified_import_name("functools", "cache")
+                    {
                         diagnostic.set_fix(Edit::replacement(
-                            unparse_expr(
+                            binding,
-                                &create_expr(ExprKind::Attribute {
-                                    value: value.clone(),
-                                    attr: "cache".to_string(),
-                                    ctx: ctx.clone(),
-                                }),
-                                checker.stylist,
-                            ),
                             expr.location,
                             expr.end_location.unwrap(),
                         ));

crates/ruff/src/rules/pyupgrade/snapshots/ruff__rules__pyupgrade__tests__UP033_0.py.snap

@@ -8,19 +8,19 @@ expression: diagnostics
     suggestion: "Rewrite with `@functools.cache"
     fixable: true
   location:
-    row: 5
+    row: 4
     column: 20
   end_location:
-    row: 5
+    row: 4
     column: 34
   fix:
     edits:
       - content: functools.cache
         location:
-          row: 5
+          row: 4
           column: 1
         end_location:
-          row: 5
+          row: 4
           column: 34
   parent: ~
 - kind:
@@ -30,32 +30,18 @@ expression: diagnostics
     fixable: true
   location:
     row: 10
-    column: 10
+    column: 20
   end_location:
     row: 10
-    column: 24
-  fix:
-    edits: []
-  parent: ~
-- kind:
-    name: LRUCacheWithMaxsizeNone
-    body: "Use `@functools.cache` instead of `@functools.lru_cache(maxsize=None)`"
-    suggestion: "Rewrite with `@functools.cache"
-    fixable: true
-  location:
-    row: 16
-    column: 20
-  end_location:
-    row: 16
     column: 34
   fix:
     edits:
       - content: functools.cache
         location:
-          row: 16
+          row: 10
           column: 1
         end_location:
-          row: 16
+          row: 10
           column: 34
   parent: ~
 - kind:
@@ -64,19 +50,19 @@ expression: diagnostics
     suggestion: "Rewrite with `@functools.cache"
     fixable: true
   location:
-    row: 21
+    row: 15
     column: 20
   end_location:
-    row: 21
+    row: 15
     column: 34
   fix:
     edits:
       - content: functools.cache
         location:
-          row: 21
+          row: 15
           column: 1
         end_location:
-          row: 21
+          row: 15
           column: 34
   parent: ~
 

crates/ruff/src/rules/pyupgrade/snapshots/ruff__rules__pyupgrade__tests__UP033_1.py.snap

@@ -0,0 +1,47 @@
+---
+source: crates/ruff/src/rules/pyupgrade/mod.rs
+expression: diagnostics
+---
+- kind:
+    name: LRUCacheWithMaxsizeNone
+    body: "Use `@functools.cache` instead of `@functools.lru_cache(maxsize=None)`"
+    suggestion: "Rewrite with `@functools.cache"
+    fixable: true
+  location:
+    row: 4
+    column: 10
+  end_location:
+    row: 4
+    column: 24
+  fix:
+    edits: []
+  parent: ~
+- kind:
+    name: LRUCacheWithMaxsizeNone
+    body: "Use `@functools.cache` instead of `@functools.lru_cache(maxsize=None)`"
+    suggestion: "Rewrite with `@functools.cache"
+    fixable: true
+  location:
+    row: 10
+    column: 10
+  end_location:
+    row: 10
+    column: 24
+  fix:
+    edits: []
+  parent: ~
+- kind:
+    name: LRUCacheWithMaxsizeNone
+    body: "Use `@functools.cache` instead of `@functools.lru_cache(maxsize=None)`"
+    suggestion: "Rewrite with `@functools.cache"
+    fixable: true
+  location:
+    row: 15
+    column: 10
+  end_location:
+    row: 15
+    column: 24
+  fix:
+    edits: []
+  parent: ~
+

crates/ruff_python_ast/src/context.rs

@@ -138,14 +138,18 @@ impl<'a> Context<'a> {
             .map_or(false, |binding| binding.kind.is_builtin())
     }
 
-    /// Resolves the call path, e.g. if you have a file
+    /// Resolves the [`Expr`] to a fully-qualified symbol-name, if `value` resolves to an imported
+    /// or builtin symbol.
+    ///
+    /// E.g., given:
+    ///
     ///
     /// ```python
     /// from sys import version_info as python_version
     /// print(python_version)
     /// ```
     ///
-    /// then `python_version` from the print statement will resolve to `sys.version_info`.
+    /// ...then `resolve_call_path(${python_version})` will resolve to `sys.version_info`.
     pub fn resolve_call_path<'b>(&'a self, value: &'b Expr) -> Option<CallPath<'a>>
     where
         'b: 'a,
@@ -203,6 +207,76 @@ impl<'a> Context<'a> {
         }
     }
 
+    /// Given a `module` and `member`, return the fully-qualified name of the binding in the current
+    /// scope, if it exists.
+    ///
+    /// E.g., given:
+    ///
+    /// ```python
+    /// from sys import version_info as python_version
+    /// print(python_version)
+    /// ```
+    ///
+    /// ...then `resolve_qualified_import_name("sys", "version_info")` will return
+    /// `Some("python_version")`.
+    pub fn resolve_qualified_import_name(&self, module: &str, member: &str) -> Option<String> {
+        self.scopes().enumerate().find_map(|(scope_index, scope)| {
+            scope.binding_ids().find_map(|binding_index| {
+                match &self.bindings[*binding_index].kind {
+                    // Ex) Given `module="sys"` and `object="exit"`:
+                    // `import sys`         -> `sys.exit`
+                    // `import sys as sys2` -> `sys2.exit`
+                    BindingKind::Importation(Importation { name, full_name }) => {
+                        if full_name == &module {
+                            // Verify that `sys` isn't bound in an inner scope.
+                            if self
+                                .scopes()
+                                .take(scope_index)
+                                .all(|scope| scope.get(name).is_none())
+                            {
+                                return Some(format!("{name}.{member}"));
+                            }
+                        }
+                    }
+                    // Ex) Given `module="os.path"` and `object="join"`:
+                    // `from os.path import join`          -> `join`
+                    // `from os.path import join as join2` -> `join2`
+                    BindingKind::FromImportation(FromImportation { name, full_name }) => {
+                        if let Some((target_module, target_member)) = full_name.split_once('.') {
+                            if target_module == module && target_member == member {
+                                // Verify that `join` isn't bound in an inner scope.
+                                if self
+                                    .scopes()
+                                    .take(scope_index)
+                                    .all(|scope| scope.get(name).is_none())
+                                {
+                                    return Some((*name).to_string());
+                                }
+                            }
+                        }
+                    }
+                    // Ex) Given `module="os"` and `object="name"`:
+                    // `import os.path ` -> `os.name`
+                    BindingKind::SubmoduleImportation(SubmoduleImportation { name, .. }) => {
+                        if name == &module {
+                            // Verify that `os` isn't bound in an inner scope.
+                            if self
+                                .scopes()
+                                .take(scope_index)
+                                .all(|scope| scope.get(name).is_none())
+                            {
+                                return Some(format!("{name}.{member}"));
+                            }
+                        }
+                    }
+                    // Non-imports.
+                    _ => {}
+                }
+                None
+            })
+        })
+    }
+
     pub fn push_parent(&mut self, parent: &'a Stmt) {
         let num_existing = self.parents.len();
         self.parents.push(RefEquality(parent));