From aea925a898d1a67170e96ebc098eb9e6325c049f Mon Sep 17 00:00:00 2001
From: Jonathan Plasse <13716151+JonathanPlasse@users.noreply.github.com>
Date: Thu, 23 Mar 2023 22:14:56 +0100
Subject: [PATCH] Fix SIM118 auto-fix (#3695)

---
 .../test/fixtures/flake8_simplify/SIM118.py   |  2 +
 .../flake8_simplify/rules/key_in_dict.rs      | 43 ++++++++++++++++---
 ...ke8_simplify__tests__SIM118_SIM118.py.snap | 20 +++++++++
 3 files changed, 60 insertions(+), 5 deletions(-)

diff --git a/crates/ruff/resources/test/fixtures/flake8_simplify/SIM118.py b/crates/ruff/resources/test/fixtures/flake8_simplify/SIM118.py
index d328992586..eaee46e84a 100644
--- a/crates/ruff/resources/test/fixtures/flake8_simplify/SIM118.py
+++ b/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
diff --git a/crates/ruff/src/rules/flake8_simplify/rules/key_in_dict.rs b/crates/ruff/src/rules/flake8_simplify/rules/key_in_dict.rs
index 6cdd2aa473..9b20274256 100644
--- a/crates/ruff/src/rules/flake8_simplify/rules/key_in_dict.rs
+++ b/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(),
         ));
diff --git a/crates/ruff/src/rules/flake8_simplify/snapshots/ruff__rules__flake8_simplify__tests__SIM118_SIM118.py.snap b/crates/ruff/src/rules/flake8_simplify/snapshots/ruff__rules__flake8_simplify__tests__SIM118_SIM118.py.snap
index 0986875eeb..c7029bcbc0 100644
--- a/crates/ruff/src/rules/flake8_simplify/snapshots/ruff__rules__flake8_simplify__tests__SIM118_SIM118.py.snap
+++ b/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: ~