Supports more cases in SIM112 (#3876)

astral-sh · Apr 4, 2023 · 390d7dc · 390d7dc
1 parent 251340a
commit 390d7dc
Show file tree

Hide file tree

Showing 3 changed files with 121 additions and 7 deletions.
diff --git 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
@@ -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/...ules/flake8_simplify/snapshots/ruff__rules__flake8_simplify__tests__SIM112_SIM112.py.snap b/...ules/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: ~