Extend dict-get-with-none-default (SIM910) to non-literals

astral-sh · Nov 18, 2023 · a8c12a4 · a8c12a4
1 parent a7fc785
commit a8c12a4
Show file tree

Hide file tree

Showing 3 changed files with 53 additions and 4 deletions.
diff --git a/crates/ruff_linter/resources/test/fixtures/flake8_simplify/SIM910.py b/crates/ruff_linter/resources/test/fixtures/flake8_simplify/SIM910.py
@@ -25,3 +25,11 @@
 
 # SIM910
 ({}).get(key, None)
+
+# SIM910
+ages = {"Tom": 23, "Maria": 23, "Dog": 11}
+age = ages.get("Cat", None)
+
+# OK
+ages = ["Tom", "Maria", "Dog"]
+age = ages.get("Cat", None)
diff --git a/crates/ruff_linter/src/rules/flake8_simplify/rules/ast_expr.rs b/crates/ruff_linter/src/rules/flake8_simplify/rules/ast_expr.rs
@@ -4,6 +4,7 @@ use ruff_text_size::Ranged;
 use crate::fix::snippet::SourceCodeSnippet;
 use ruff_diagnostics::{AlwaysFixableViolation, Diagnostic, Edit, Fix, FixAvailability, Violation};
 use ruff_macros::{derive_message_formats, violation};
+use ruff_python_semantic::analyze::typing::is_dict;
 
 use crate::checkers::ast::Checker;
 
@@ -244,9 +245,6 @@ pub(crate) fn dict_get_with_none_default(checker: &mut Checker, expr: &Expr) {
     let Expr::Attribute(ast::ExprAttribute { value, attr, .. }) = func.as_ref() else {
         return;
     };
-    if !value.is_dict_expr() {
-        return;
-    }
     if attr != "get" {
         return;
     }
@@ -263,6 +261,24 @@ pub(crate) fn dict_get_with_none_default(checker: &mut Checker, expr: &Expr) {
         return;
     }
 
+    // Check if the value is a dictionary.
+    match value.as_ref() {
+        Expr::Dict(_) | Expr::DictComp(_) => {}
+        Expr::Name(name) => {
+            let Some(binding) = checker
+                .semantic()
+                .only_binding(name)
+                .map(|id| checker.semantic().binding(id))
+            else {
+                return;
+            };
+            if !is_dict(binding, checker.semantic()) {
+                return;
+            }
+        }
+        _ => return,
+    }
+
     let expected = format!(
         "{}({})",
         checker.locator().slice(func.as_ref()),
@@ -277,7 +293,6 @@ pub(crate) fn dict_get_with_none_default(checker: &mut Checker, expr: &Expr) {
         },
         expr.range(),
     );
-
     diagnostic.set_fix(Fix::safe_edit(Edit::range_replacement(
         expected,
         expr.range(),

diff --git a/...ake8_simplify/snapshots/ruff_linter__rules__flake8_simplify__tests__SIM910_SIM910.py.snap b/...ake8_simplify/snapshots/ruff_linter__rules__flake8_simplify__tests__SIM910_SIM910.py.snap
@@ -83,6 +83,8 @@ SIM910.py:27:1: SIM910 [*] Use `({}).get(key)` instead of `({}).get(key, None)`
 26 | # SIM910
 27 | ({}).get(key, None)
    | ^^^^^^^^^^^^^^^^^^^ SIM910
+28 | 
+29 | # SIM910
    |
    = help: Replace `({}).get(key, None)` with `({}).get(key)`
 
@@ -92,5 +94,29 @@ SIM910.py:27:1: SIM910 [*] Use `({}).get(key)` instead of `({}).get(key, None)`
 26 26 | # SIM910
 27    |-({}).get(key, None)
    27 |+({}).get(key)
+28 28 | 
+29 29 | # SIM910
+30 30 | ages = {"Tom": 23, "Maria": 23, "Dog": 11}
+
+SIM910.py:31:7: SIM910 [*] Use `ages.get("Cat")` instead of `ages.get("Cat", None)`
+   |
+29 | # SIM910
+30 | ages = {"Tom": 23, "Maria": 23, "Dog": 11}
+31 | age = ages.get("Cat", None)
+   |       ^^^^^^^^^^^^^^^^^^^^^ SIM910
+32 | 
+33 | # OK
+   |
+   = help: Replace `ages.get("Cat", None)` with `ages.get("Cat")`
+
+ℹ Safe fix
+28 28 | 
+29 29 | # SIM910
+30 30 | ages = {"Tom": 23, "Maria": 23, "Dog": 11}
+31    |-age = ages.get("Cat", None)
+   31 |+age = ages.get("Cat")
+32 32 | 
+33 33 | # OK
+34 34 | ages = ["Tom", "Maria", "Dog"]