Implement flake8-simplify SIM112 (#1764)

Ref #998

README.md

@@ -980,6 +980,7 @@ For more, see [flake8-simplify](https://pypi.org/project/flake8-simplify/0.19.3/
 | SIM109 | CompareWithTuple | Use `value in (..., ...)` instead of `value == ... or value == ...` | 🛠 |
 | SIM110 | ConvertLoopToAny | Use `return any(x for x in y)` instead of `for` loop | 🛠 |
 | SIM111 | ConvertLoopToAll | Use `return all(x for x in y)` instead of `for` loop | 🛠 |
+| SIM112 | UseCapitalEnvironmentVariables | Use capitalized environment variable `...` instead of `...` | 🛠 |
 | SIM117 | MultipleWithStatements | Use a single `with` statement with multiple contexts instead of nested `with` statements |  |
 | SIM118 | KeyInDict | Use `key in dict` instead of `key in dict.keys()` | 🛠 |
 | SIM201 | NegateEqualOp | Use `left != right` instead of `not left == right` | 🛠 |

resources/test/fixtures/flake8_simplify/SIM112.py

@@ -0,0 +1,19 @@
+import os
+
+# Bad
+os.environ['foo']
+
+os.environ.get('foo')
+
+os.environ.get('foo', 'bar')
+
+os.getenv('foo')
+
+# Good
+os.environ['FOO']
+
+os.environ.get('FOO')
+
+os.environ.get('FOO', 'bar')
+
+os.getenv('FOO')

ruff.schema.json

@@ -1506,6 +1506,7 @@
         "SIM11",
         "SIM110",
         "SIM111",
+        "SIM112",
         "SIM117",
         "SIM118",
         "SIM2",

src/checkers/ast.rs

@@ -1405,6 +1405,9 @@ where
                 if self.settings.enabled.contains(&RuleCode::B015) {
                     flake8_bugbear::rules::useless_comparison(self, value);
                 }
+                if self.settings.enabled.contains(&RuleCode::SIM112) {
+                    flake8_simplify::rules::use_capital_environment_variables(self, value);
+                }
             }
             _ => {}
         }

src/flake8_simplify/mod.rs

@@ -21,6 +21,7 @@ mod tests {
     #[test_case(RuleCode::SIM109, Path::new("SIM109.py"); "SIM109")]
     #[test_case(RuleCode::SIM110, Path::new("SIM110.py"); "SIM110")]
     #[test_case(RuleCode::SIM111, Path::new("SIM111.py"); "SIM111")]
+    #[test_case(RuleCode::SIM112, Path::new("SIM112.py"); "SIM112")]
     #[test_case(RuleCode::SIM117, Path::new("SIM117.py"); "SIM117")]
     #[test_case(RuleCode::SIM201, Path::new("SIM201.py"); "SIM201")]
     #[test_case(RuleCode::SIM202, Path::new("SIM202.py"); "SIM202")]

src/flake8_simplify/rules/ast_expr.rs

@@ -0,0 +1,106 @@
+use rustpython_ast::{Constant, Expr, ExprKind};
+
+use crate::ast::helpers::{create_expr, match_module_member, unparse_expr};
+use crate::ast::types::Range;
+use crate::autofix::Fix;
+use crate::checkers::ast::Checker;
+use crate::registry::{Diagnostic, RuleCode};
+use crate::violations;
+
+/// SIM112
+pub fn use_capital_environment_variables(checker: &mut Checker, expr: &Expr) {
+    // check `os.environ['foo']`
+    if let ExprKind::Subscript { .. } = &expr.node {
+        check_os_environ_subscript(checker, expr);
+        return;
+    }
+
+    // check `os.environ.get('foo')` and `os.getenv('foo')``
+    let is_os_environ_get = match_module_member(
+        expr,
+        "os.environ",
+        "get",
+        &checker.from_imports,
+        &checker.import_aliases,
+    );
+    let is_os_getenv = match_module_member(
+        expr,
+        "os",
+        "getenv",
+        &checker.from_imports,
+        &checker.import_aliases,
+    );
+    if !(is_os_environ_get || is_os_getenv) {
+        return;
+    }
+
+    let ExprKind::Call { args, .. } = &expr.node else {
+        return;
+    };
+    let Some(arg) = args.get(0) else {
+        return;
+    };
+    let ExprKind::Constant { value: Constant::Str(env_var), kind } = &arg.node else {
+        return;
+    };
+    let capital_env_var = env_var.to_ascii_uppercase();
+    if &capital_env_var == env_var {
+        return;
+    }
+
+    let mut diagnostic = Diagnostic::new(
+        violations::UseCapitalEnvironmentVariables(capital_env_var.clone(), env_var.clone()),
+        Range::from_located(arg),
+    );
+    if checker.patch(&RuleCode::SIM112) {
+        let new_env_var = create_expr(ExprKind::Constant {
+            value: capital_env_var.into(),
+            kind: kind.clone(),
+        });
+        diagnostic.amend(Fix::replacement(
+            unparse_expr(&new_env_var, checker.style),
+            arg.location,
+            arg.end_location.unwrap(),
+        ));
+    }
+    checker.diagnostics.push(diagnostic);
+}
+
+fn check_os_environ_subscript(checker: &mut Checker, expr: &Expr) {
+    let ExprKind::Subscript { value, slice, .. } = &expr.node else {
+        return;
+    };
+    let ExprKind::Attribute { value: attr_value, attr, .. } = &value.node else {
+        return;
+    };
+    let ExprKind::Name { id, .. } = &attr_value.node else {
+        return;
+    };
+    if id != "os" || attr != "environ" {
+        return;
+    }
+    let ExprKind::Constant { value: Constant::Str(env_var), kind } = &slice.node else {
+        return;
+    };
+    let capital_env_var = env_var.to_ascii_uppercase();
+    if &capital_env_var == env_var {
+        return;
+    }
+
+    let mut diagnostic = Diagnostic::new(
+        violations::UseCapitalEnvironmentVariables(capital_env_var.clone(), env_var.clone()),
+        Range::from_located(slice),
+    );
+    if checker.patch(&RuleCode::SIM112) {
+        let new_env_var = create_expr(ExprKind::Constant {
+            value: capital_env_var.into(),
+            kind: kind.clone(),
+        });
+        diagnostic.amend(Fix::replacement(
+            unparse_expr(&new_env_var, checker.style),
+            slice.location,
+            slice.end_location.unwrap(),
+        ));
+    }
+    checker.diagnostics.push(diagnostic);
+}

src/flake8_simplify/rules/mod.rs

@@ -1,6 +1,7 @@
 pub use ast_bool_op::{
     a_and_not_a, a_or_not_a, and_false, compare_with_tuple, duplicate_isinstance_call, or_true,
 };
+pub use ast_expr::use_capital_environment_variables;
 pub use ast_for::convert_loop_to_any_all;
 pub use ast_if::{nested_if_statements, return_bool_condition_directly, use_ternary_operator};
 pub use ast_ifexp::{
@@ -14,6 +15,7 @@ pub use use_contextlib_suppress::use_contextlib_suppress;
 pub use yoda_conditions::yoda_conditions;
 
 mod ast_bool_op;
+mod ast_expr;
 mod ast_for;
 mod ast_if;
 mod ast_ifexp;

src/flake8_simplify/snapshots/ruff__flake8_simplify__tests__SIM112_SIM112.py.snap

@@ -0,0 +1,81 @@
+---
+source: src/flake8_simplify/mod.rs
+expression: diagnostics
+---
+- kind:
+    UseCapitalEnvironmentVariables:
+      - FOO
+      - foo
+  location:
+    row: 4
+    column: 11
+  end_location:
+    row: 4
+    column: 16
+  fix:
+    content: "'FOO'"
+    location:
+      row: 4
+      column: 11
+    end_location:
+      row: 4
+      column: 16
+  parent: ~
+- kind:
+    UseCapitalEnvironmentVariables:
+      - FOO
+      - foo
+  location:
+    row: 6
+    column: 15
+  end_location:
+    row: 6
+    column: 20
+  fix:
+    content: "'FOO'"
+    location:
+      row: 6
+      column: 15
+    end_location:
+      row: 6
+      column: 20
+  parent: ~
+- kind:
+    UseCapitalEnvironmentVariables:
+      - FOO
+      - foo
+  location:
+    row: 8
+    column: 15
+  end_location:
+    row: 8
+    column: 20
+  fix:
+    content: "'FOO'"
+    location:
+      row: 8
+      column: 15
+    end_location:
+      row: 8
+      column: 20
+  parent: ~
+- kind:
+    UseCapitalEnvironmentVariables:
+      - FOO
+      - foo
+  location:
+    row: 10
+    column: 10
+  end_location:
+    row: 10
+    column: 15
+  fix:
+    content: "'FOO'"
+    location:
+      row: 10
+      column: 10
+    end_location:
+      row: 10
+      column: 15
+  parent: ~
+

src/registry.rs

@@ -295,6 +295,7 @@ define_rule_mapping!(
     SIM109 => violations::CompareWithTuple,
     SIM110 => violations::ConvertLoopToAny,
     SIM111 => violations::ConvertLoopToAll,
+    SIM112 => violations::UseCapitalEnvironmentVariables,
     SIM117 => violations::MultipleWithStatements,
     SIM118 => violations::KeyInDict,
     SIM201 => violations::NegateEqualOp,
@@ -1110,6 +1111,7 @@ impl RuleCode {
             RuleCode::SIM109 => RuleOrigin::Flake8Simplify,
             RuleCode::SIM110 => RuleOrigin::Flake8Simplify,
             RuleCode::SIM111 => RuleOrigin::Flake8Simplify,
+            RuleCode::SIM112 => RuleOrigin::Flake8Simplify,
             RuleCode::SIM117 => RuleOrigin::Flake8Simplify,
             RuleCode::SIM118 => RuleOrigin::Flake8Simplify,
             RuleCode::SIM201 => RuleOrigin::Flake8Simplify,

src/violations.rs

@@ -2677,6 +2677,24 @@ impl Violation for SysVersionSlice1Referenced {
 }
 
 // flake8-simplify
+define_violation!(
+    pub struct UseCapitalEnvironmentVariables(pub String, pub String);
+);
+impl AlwaysAutofixableViolation for UseCapitalEnvironmentVariables {
+    fn message(&self) -> String {
+        let UseCapitalEnvironmentVariables(expected, original) = self;
+        format!("Use capitalized environment variable `{expected}` instead of `{original}`")
+    }
+
+    fn autofix_title(&self) -> String {
+        let UseCapitalEnvironmentVariables(expected, original) = self;
+        format!("Replace `{original}` with `{expected}`")
+    }
+
+    fn placeholder() -> Self {
+        UseCapitalEnvironmentVariables("...".to_string(), "...".to_string())
+    }
+}
 
 define_violation!(
     pub struct DuplicateIsinstanceCall(pub String);