Rewrite lru_cache to cache on Python 3.9+ (#1918)

Closes #1913.
astral-sh · Jan 16, 2023 · c0845a8 · c0845a8
1 parent 019ecc4
commit c0845a8
Show file tree

Hide file tree

Showing 2 changed files with 60 additions and 59 deletions.
diff --git a/src/rules/pyupgrade/rules/unnecessary_lru_cache_params.rs b/src/rules/pyupgrade/rules/unnecessary_lru_cache_params.rs
@@ -1,28 +1,26 @@
 use rustpython_ast::{Constant, ExprKind, KeywordData};
 use rustpython_parser::ast::Expr;
 
+use crate::ast::helpers::{create_expr, unparse_expr};
 use crate::ast::types::Range;
 use crate::checkers::ast::Checker;
 use crate::fix::Fix;
-use crate::registry::Diagnostic;
+use crate::registry::{Diagnostic, RuleCode};
 use crate::settings::types::PythonVersion;
 use crate::violations;
 
-fn rule(
-    checker: &Checker,
-    decorator_list: &[Expr],
-    target_version: PythonVersion,
-) -> Option<Diagnostic> {
+/// UP011
+pub fn unnecessary_lru_cache_params(checker: &mut Checker, decorator_list: &[Expr]) {
     for expr in decorator_list.iter() {
         let ExprKind::Call {
             func,
             args,
             keywords,
-        } = &expr.node
-        else {
+        } = &expr.node else {
             continue;
         };
 
+        // Look for, e.g., `import functools; @functools.lru_cache`.
         if !(args.is_empty()
             && checker
                 .resolve_call_path(func)
@@ -31,21 +29,29 @@ fn rule(
             continue;
         }
 
-        let range = Range::new(func.end_location.unwrap(), expr.end_location.unwrap());
         // Ex) `functools.lru_cache()`
         if keywords.is_empty() {
-            return Some(Diagnostic::new(
+            let mut diagnostic = Diagnostic::new(
                 violations::UnnecessaryLRUCacheParams,
-                range,
-            ));
+                Range::new(func.end_location.unwrap(), expr.end_location.unwrap()),
+            );
+            if checker.patch(&RuleCode::UP011) {
+                diagnostic.amend(Fix::replacement(
+                    unparse_expr(func, checker.stylist),
+                    expr.location,
+                    expr.end_location.unwrap(),
+                ));
+            }
+            checker.diagnostics.push(diagnostic);
         }
+
         // Ex) `functools.lru_cache(maxsize=None)`
-        if !(target_version >= PythonVersion::Py39 && keywords.len() == 1) {
+        if !(checker.settings.target_version >= PythonVersion::Py39 && keywords.len() == 1) {
             continue;
         }
 
         let KeywordData { arg, value } = &keywords[0].node;
-        if !(arg.as_ref().map(|arg| arg == "maxsize").unwrap_or_default()
+        if !(arg.as_ref().map_or(false, |arg| arg == "maxsize")
             && matches!(
                 value.node,
                 ExprKind::Constant {
@@ -56,25 +62,27 @@ fn rule(
         {
             continue;
         }
-        return Some(Diagnostic::new(
-            violations::UnnecessaryLRUCacheParams,
-            range,
-        ));
-    }
-    None
-}
 
-/// UP011
-pub fn unnecessary_lru_cache_params(checker: &mut Checker, decorator_list: &[Expr]) {
-    let Some(mut diagnostic) = rule(
-        checker,
-        decorator_list,
-        checker.settings.target_version,
-    ) else {
-        return;
-    };
-    if checker.patch(diagnostic.kind.code()) {
-        diagnostic.amend(Fix::deletion(diagnostic.location, diagnostic.end_location));
+        let mut diagnostic = Diagnostic::new(
+            violations::UnnecessaryLRUCacheParams,
+            Range::new(func.end_location.unwrap(), expr.end_location.unwrap()),
+        );
+        if checker.patch(&RuleCode::UP011) {
+            if let ExprKind::Attribute { value, ctx, .. } = &func.node {
+                diagnostic.amend(Fix::replacement(
+                    unparse_expr(
+                        &create_expr(ExprKind::Attribute {
+                            value: value.clone(),
+                            attr: "cache".to_string(),
+                            ctx: ctx.clone(),
+                        }),
+                        checker.stylist,
+                    ),
+                    expr.location,
+                    expr.end_location.unwrap(),
+                ));
+            }
+        }
+        checker.diagnostics.push(diagnostic);
     }
-    checker.diagnostics.push(diagnostic);
 }
diff --git a/src/rules/pyupgrade/snapshots/ruff__rules__pyupgrade__tests__UP011_UP011_0.py.snap b/src/rules/pyupgrade/snapshots/ruff__rules__pyupgrade__tests__UP011_UP011_0.py.snap
@@ -11,10 +11,10 @@ expression: diagnostics
     row: 5
     column: 12
   fix:
-    content: ""
+    content: lru_cache
     location:
       row: 5
-      column: 10
+      column: 1
     end_location:
       row: 5
       column: 12
@@ -28,10 +28,10 @@ expression: diagnostics
     row: 11
     column: 22
   fix:
-    content: ""
+    content: functools.lru_cache
     location:
       row: 11
-      column: 20
+      column: 1
     end_location:
       row: 11
       column: 22
@@ -44,14 +44,7 @@ expression: diagnostics
   end_location:
     row: 16
     column: 24
-  fix:
-    content: ""
-    location:
-      row: 16
-      column: 10
-    end_location:
-      row: 16
-      column: 24
+  fix: ~
   parent: ~
 - kind:
     UnnecessaryLRUCacheParams: ~
@@ -62,10 +55,10 @@ expression: diagnostics
     row: 21
     column: 34
   fix:
-    content: ""
+    content: functools.cache
     location:
       row: 21
-      column: 20
+      column: 1
     end_location:
       row: 21
       column: 34
@@ -79,10 +72,10 @@ expression: diagnostics
     row: 28
     column: 1
   fix:
-    content: ""
+    content: lru_cache
     location:
       row: 27
-      column: 10
+      column: 1
     end_location:
       row: 28
       column: 1
@@ -96,10 +89,10 @@ expression: diagnostics
     row: 35
     column: 1
   fix:
-    content: ""
+    content: lru_cache
     location:
       row: 33
-      column: 10
+      column: 1
     end_location:
       row: 35
       column: 1
@@ -113,10 +106,10 @@ expression: diagnostics
     row: 42
     column: 19
   fix:
-    content: ""
+    content: functools.cache
     location:
       row: 40
-      column: 20
+      column: 1
     end_location:
       row: 42
       column: 19
@@ -130,10 +123,10 @@ expression: diagnostics
     row: 51
     column: 1
   fix:
-    content: ""
+    content: functools.cache
     location:
       row: 47
-      column: 20
+      column: 1
     end_location:
       row: 51
       column: 1
@@ -147,10 +140,10 @@ expression: diagnostics
     row: 62
     column: 1
   fix:
-    content: ""
+    content: functools.cache
     location:
       row: 56
-      column: 20
+      column: 1
     end_location:
       row: 62
       column: 1
@@ -164,10 +157,10 @@ expression: diagnostics
     row: 72
     column: 1
   fix:
-    content: ""
+    content: functools.cache
     location:
       row: 67
-      column: 20
+      column: 1
     end_location:
       row: 72
       column: 1