From 390d7dcf391cbdbe4df80e3377c295a3bc03db5d Mon Sep 17 00:00:00 2001
From: kyoto7250 <50972773+kyoto7250@users.noreply.github.com>
Date: Wed, 5 Apr 2023 04:49:24 +0900
Subject: [PATCH] Supports more cases in `SIM112` (#3876)

---
 .../test/fixtures/flake8_simplify/SIM112.py   | 21 +++++
 crates/ruff/src/checkers/ast/mod.rs           | 23 +++--
 ...ke8_simplify__tests__SIM112_SIM112.py.snap | 84 +++++++++++++++++++
 3 files changed, 121 insertions(+), 7 deletions(-)

diff --git a/crates/ruff/resources/test/fixtures/flake8_simplify/SIM112.py b/crates/ruff/resources/test/fixtures/flake8_simplify/SIM112.py
index 01cfe00320480..74b44d63325dc 100644
--- a/crates/ruff/resources/test/fixtures/flake8_simplify/SIM112.py
+++ b/crates/ruff/resources/test/fixtures/flake8_simplify/SIM112.py
@@ -9,6 +9,17 @@
 
 os.getenv('foo')
 
+env = os.environ.get('foo')
+
+env = os.environ['foo']
+
+if env := os.environ.get('foo'):
+    pass
+
+if env := os.environ['foo']:
+    pass
+
+
 # Good
 os.environ['FOO']
 
@@ -17,3 +28,13 @@
 os.environ.get('FOO', 'bar')
 
 os.getenv('FOO')
+
+env = os.getenv('FOO')
+
+if env := os.getenv('FOO'):
+    pass
+
+env = os.environ['FOO']
+
+if env := os.environ['FOO']:
+    pass
diff --git a/crates/ruff/src/checkers/ast/mod.rs b/crates/ruff/src/checkers/ast/mod.rs
index 735e5313cc839..ff33cbbbb434e 100644
--- a/crates/ruff/src/checkers/ast/mod.rs
+++ b/crates/ruff/src/checkers/ast/mod.rs
@@ -1847,13 +1847,6 @@ where
                 if self.settings.rules.enabled(Rule::UselessExpression) {
                     flake8_bugbear::rules::useless_expression(self, value);
                 }
-                if self
-                    .settings
-                    .rules
-                    .enabled(Rule::UncapitalizedEnvironmentVariables)
-                {
-                    flake8_simplify::rules::use_capital_environment_variables(self, value);
-                }
                 if self.settings.rules.enabled(Rule::AsyncioDanglingTask) {
                     if let Some(diagnostic) = ruff::rules::asyncio_dangling_task(value, |expr| {
                         self.ctx.resolve_call_path(expr)
@@ -2210,6 +2203,14 @@ where
                 ]) {
                     flake8_2020::rules::subscript(self, value, slice);
                 }
+
+                if self
+                    .settings
+                    .rules
+                    .enabled(Rule::UncapitalizedEnvironmentVariables)
+                {
+                    flake8_simplify::rules::use_capital_environment_variables(self, expr);
+                }
             }
             ExprKind::Tuple { elts, ctx } | ExprKind::List { elts, ctx } => {
                 if matches!(ctx, ExprContext::Store) {
@@ -2910,6 +2911,14 @@ where
                 }
 
                 // flake8-simplify
+                if self
+                    .settings
+                    .rules
+                    .enabled(Rule::UncapitalizedEnvironmentVariables)
+                {
+                    flake8_simplify::rules::use_capital_environment_variables(self, expr);
+                }
+
                 if self
                     .settings
                     .rules
diff --git a/crates/ruff/src/rules/flake8_simplify/snapshots/ruff__rules__flake8_simplify__tests__SIM112_SIM112.py.snap b/crates/ruff/src/rules/flake8_simplify/snapshots/ruff__rules__flake8_simplify__tests__SIM112_SIM112.py.snap
index 107d61483c201..54bea4d1f986a 100644
--- a/crates/ruff/src/rules/flake8_simplify/snapshots/ruff__rules__flake8_simplify__tests__SIM112_SIM112.py.snap
+++ b/crates/ruff/src/rules/flake8_simplify/snapshots/ruff__rules__flake8_simplify__tests__SIM112_SIM112.py.snap
@@ -86,4 +86,88 @@ expression: diagnostics
           row: 10
           column: 15
   parent: ~
+- kind:
+    name: UncapitalizedEnvironmentVariables
+    body: "Use capitalized environment variable `FOO` instead of `foo`"
+    suggestion: "Replace `foo` with `FOO`"
+    fixable: true
+  location:
+    row: 12
+    column: 21
+  end_location:
+    row: 12
+    column: 26
+  fix:
+    edits:
+      - content: "'FOO'"
+        location:
+          row: 12
+          column: 21
+        end_location:
+          row: 12
+          column: 26
+  parent: ~
+- kind:
+    name: UncapitalizedEnvironmentVariables
+    body: "Use capitalized environment variable `FOO` instead of `foo`"
+    suggestion: "Replace `foo` with `FOO`"
+    fixable: true
+  location:
+    row: 14
+    column: 17
+  end_location:
+    row: 14
+    column: 22
+  fix:
+    edits:
+      - content: "'FOO'"
+        location:
+          row: 14
+          column: 17
+        end_location:
+          row: 14
+          column: 22
+  parent: ~
+- kind:
+    name: UncapitalizedEnvironmentVariables
+    body: "Use capitalized environment variable `FOO` instead of `foo`"
+    suggestion: "Replace `foo` with `FOO`"
+    fixable: true
+  location:
+    row: 16
+    column: 25
+  end_location:
+    row: 16
+    column: 30
+  fix:
+    edits:
+      - content: "'FOO'"
+        location:
+          row: 16
+          column: 25
+        end_location:
+          row: 16
+          column: 30
+  parent: ~
+- kind:
+    name: UncapitalizedEnvironmentVariables
+    body: "Use capitalized environment variable `FOO` instead of `foo`"
+    suggestion: "Replace `foo` with `FOO`"
+    fixable: true
+  location:
+    row: 19
+    column: 21
+  end_location:
+    row: 19
+    column: 26
+  fix:
+    edits:
+      - content: "'FOO'"
+        location:
+          row: 19
+          column: 21
+        end_location:
+          row: 19
+          column: 26
+  parent: ~