Fix SIM118 auto-fix (#3695)

crates/ruff/resources/test/fixtures/flake8_simplify/SIM118.py

@@ -20,3 +20,5 @@ for key in list(obj.keys()):
 {k: k for k in obj.keys()}  # SIM118
 
 (k for k in obj.keys())  # SIM118
+
+key in (obj or {}).keys()  # SIM118

crates/ruff/src/rules/flake8_simplify/rules/key_in_dict.rs

@@ -1,10 +1,16 @@
+use anyhow::Result;
+use libcst_native::{Codegen, CodegenState};
+use log::error;
 use rustpython_parser::ast::{Cmpop, Expr, ExprKind};
 
-use ruff_diagnostics::{AlwaysAutofixableViolation, Diagnostic, Fix};
+use ruff_diagnostics::Fix;
+use ruff_diagnostics::{AlwaysAutofixableViolation, Diagnostic};
 use ruff_macros::{derive_message_formats, violation};
+use ruff_python_ast::source_code::{Locator, Stylist};
 use ruff_python_ast::types::Range;
 
 use crate::checkers::ast::Checker;
+use crate::cst::matchers::{match_attribute, match_call, match_expression};
 use crate::registry::AsRule;
 
 #[violation]
@@ -26,6 +32,26 @@ impl AlwaysAutofixableViolation for InDictKeys {
     }
 }
 
+fn get_value_content_for_key_in_dict(
+    locator: &Locator,
+    stylist: &Stylist,
+    expr: &rustpython_parser::ast::Expr,
+) -> Result<String> {
+    let content = locator.slice(expr);
+    let mut expression = match_expression(content)?;
+    let call = match_call(&mut expression)?;
+    let attribute = match_attribute(&mut call.func)?;
+
+    let mut state = CodegenState {
+        default_newline: stylist.line_ending(),
+        default_indent: stylist.indentation(),
+        ..CodegenState::default()
+    };
+    attribute.value.codegen(&mut state);
+
+    Ok(state.to_string())
+}
+
 /// SIM118
 fn key_in_dict(checker: &mut Checker, left: &Expr, right: &Expr, range: Range) {
     let ExprKind::Call {
@@ -39,7 +65,7 @@ fn key_in_dict(checker: &mut Checker, left: &Expr, right: &Expr, range: Range) {
         return;
     }
 
-    let ExprKind::Attribute { attr, value, .. } = &func.node else {
+    let ExprKind::Attribute { attr, .. } = &func.node else {
         return;
     };
     if attr != "keys" {
@@ -48,18 +74,25 @@ fn key_in_dict(checker: &mut Checker, left: &Expr, right: &Expr, range: Range) {
 
     // Slice exact content to preserve formatting.
     let left_content = checker.locator.slice(left);
-    let value_content = checker.locator.slice(value);
+    let value_content =
+        match get_value_content_for_key_in_dict(checker.locator, checker.stylist, right) {
+            Ok(value_content) => value_content,
+            Err(err) => {
+                error!("Failed to get value content for key in dict: {}", err);
+                return;
+            }
+        };
 
     let mut diagnostic = Diagnostic::new(
         InDictKeys {
             key: left_content.to_string(),
-            dict: value_content.to_string(),
+            dict: value_content.clone(),
         },
         range,
     );
     if checker.patch(diagnostic.kind.rule()) {
         diagnostic.amend(Fix::replacement(
-            value_content.to_string(),
+            value_content,
             right.location,
             right.end_location.unwrap(),
         ));

crates/ruff/src/rules/flake8_simplify/snapshots/ruff__rules__flake8_simplify__tests__SIM118_SIM118.py.snap

@@ -182,4 +182,24 @@ expression: diagnostics
       row: 22
       column: 22
   parent: ~
+- kind:
+    name: InDictKeys
+    body: "Use `key in (obj or {})` instead of `key in (obj or {}).keys()`"
+    suggestion: "Convert to `key in (obj or {})`"
+    fixable: true
+  location:
+    row: 24
+    column: 0
+  end_location:
+    row: 24
+    column: 25
+  fix:
+    content: "(obj or {})"
+    location:
+      row: 24
+      column: 7
+    end_location:
+      row: 24
+      column: 25
+  parent: ~