Allow f-strings and concatenations in os.getenv (#3516)

crates/ruff/resources/test/fixtures/pylint/invalid_envvar_default.py

@@ -5,4 +5,8 @@ goodVar = os.getenv("TESTING", None)
 dictVarBad = os.getenv("AAA", {"a", 7})  # [invalid-envvar-default]
 print(os.getenv("TEST", False))  # [invalid-envvar-default]
 os.getenv("AA", "GOOD")
+os.getenv("AA", f"GOOD")
+os.getenv("AA", "GOOD" + "BAD")
+os.getenv("AA", "GOOD" + 1)
 os.getenv("B", Z)
+

crates/ruff/resources/test/fixtures/pylint/invalid_envvar_value.py

@@ -7,6 +7,9 @@ os.getenv(key="testingAgain")
 os.getenv(key=11)  # [invalid-envvar-value]
 os.getenv(["hello"])  # [invalid-envvar-value]
 os.getenv(key="foo", default="bar")
+os.getenv(key=f"foo", default="bar")
+os.getenv(key="foo" + "bar", default=1)
+os.getenv(key=1 + "bar", default=1)  # [invalid-envvar-value]
 
 AA = "aa"
 os.getenv(AA)

crates/ruff/src/rules/pylint/rules/invalid_envvar_default.rs

@@ -1,4 +1,4 @@
-use rustpython_parser::ast::{Constant, Expr, ExprKind, Keyword};
+use rustpython_parser::ast::{Constant, Expr, ExprKind, Keyword, Operator};
 
 use ruff_diagnostics::{Diagnostic, Violation};
 use ruff_macros::{derive_message_formats, violation};
@@ -49,13 +49,22 @@ fn is_valid_default(expr: &Expr) -> bool {
         return true;
     }
 
+    if let ExprKind::BinOp {
+        left,
+        right,
+        op: Operator::Add,
+    } = &expr.node
+    {
+        return is_valid_default(left) && is_valid_default(right);
+    }
+
     // Otherwise, the default must be a string or `None`.
     matches!(
         expr.node,
         ExprKind::Constant {
             value: Constant::Str { .. } | Constant::None { .. },
             ..
-        }
+        } | ExprKind::JoinedStr { .. }
     )
 }
 

crates/ruff/src/rules/pylint/rules/invalid_envvar_value.rs

@@ -1,4 +1,4 @@
-use rustpython_parser::ast::{Constant, Expr, ExprKind, Keyword};
+use rustpython_parser::ast::{Constant, Expr, ExprKind, Keyword, Operator};
 
 use ruff_diagnostics::{Diagnostic, Violation};
 use ruff_macros::{derive_message_formats, violation};
@@ -34,6 +34,37 @@ impl Violation for InvalidEnvvarValue {
     }
 }
 
+fn is_valid_key(expr: &Expr) -> bool {
+    // We can't infer the types of these defaults, so assume they're valid.
+    if matches!(
+        expr.node,
+        ExprKind::Name { .. }
+            | ExprKind::Attribute { .. }
+            | ExprKind::Subscript { .. }
+            | ExprKind::Call { .. }
+    ) {
+        return true;
+    }
+
+    if let ExprKind::BinOp {
+        left,
+        right,
+        op: Operator::Add,
+    } = &expr.node
+    {
+        return is_valid_key(left) && is_valid_key(right);
+    }
+
+    // Otherwise, the default must be a string.
+    matches!(
+        expr.node,
+        ExprKind::Constant {
+            value: Constant::Str { .. },
+            ..
+        } | ExprKind::JoinedStr { .. }
+    )
+}
+
 /// PLE1507
 pub fn invalid_envvar_value(
     checker: &mut Checker,
@@ -46,28 +77,20 @@ pub fn invalid_envvar_value(
         .resolve_call_path(func)
         .map_or(false, |call_path| call_path.as_slice() == ["os", "getenv"])
     {
-        // Get the first argument for `getenv`
-        if let Some(expr) = args.get(0).or_else(|| {
+        // Find the `key` argument, if it exists.
+        let Some(expr) = args.get(0).or_else(|| {
             keywords
                 .iter()
                 .find(|keyword| keyword.node.arg.as_ref().map_or(false, |arg| arg == "key"))
                 .map(|keyword| &keyword.node.value)
-        }) {
-            // Ignoring types that are inferred, only do direct constants
-            if !matches!(
-                expr.node,
-                ExprKind::Constant {
-                    value: Constant::Str { .. },
-                    ..
-                } | ExprKind::Name { .. }
-                    | ExprKind::Attribute { .. }
-                    | ExprKind::Subscript { .. }
-                    | ExprKind::Call { .. }
-            ) {
-                checker
-                    .diagnostics
-                    .push(Diagnostic::new(InvalidEnvvarValue, Range::from(expr)));
-            }
+        }) else {
+            return;
+        };
+
+        if !is_valid_key(expr) {
+            checker
+                .diagnostics
+                .push(Diagnostic::new(InvalidEnvvarValue, Range::from(expr)));
         }
     }
 }

crates/ruff/src/rules/pylint/snapshots/ruff__rules__pylint__tests__PLE1507_invalid_envvar_value.py.snap

@@ -41,4 +41,17 @@ expression: diagnostics
     column: 19
   fix: ~
   parent: ~
+- kind:
+    name: InvalidEnvvarValue
+    body: "Invalid type for initial `os.getenv` argument; expected `str`"
+    suggestion: ~
+    fixable: false
+  location:
+    row: 12
+    column: 14
+  end_location:
+    row: 12
+    column: 23
+  fix: ~
+  parent: ~
 

crates/ruff/src/rules/pylint/snapshots/ruff__rules__pylint__tests__PLW1508_invalid_envvar_default.py.snap

@@ -41,4 +41,17 @@ expression: diagnostics
     column: 29
   fix: ~
   parent: ~
+- kind:
+    name: InvalidEnvvarDefault
+    body: "Invalid type for environment variable default; expected `str` or `None`"
+    suggestion: ~
+    fixable: false
+  location:
+    row: 10
+    column: 16
+  end_location:
+    row: 10
+    column: 26
+  fix: ~
+  parent: ~