From 7ae7bf6e3060eb73ff9be95216f7a3384161de27 Mon Sep 17 00:00:00 2001
From: Christopher Covington <cov@mykolab.com>
Date: Wed, 31 Jan 2024 15:41:24 +0000
Subject: [PATCH] Support `IfExp` with dual string arms in
 `invalid-envvar-default` (#9734)

## Summary

Just like #6537 and #6538 but for the `default` second parameter to
`getenv()`.

Also rename "BAD" to "BAR" in the tests, since those strings shouldn't
trigger the rule.

## Test Plan

Added passing and failing examples to `invalid_envvar_default.py`.
---
 .../fixtures/pylint/invalid_envvar_default.py |  7 +--
 .../pylint/rules/invalid_envvar_default.rs    | 53 ++++---------------
 ...ts__PLW1508_invalid_envvar_default.py.snap | 12 ++++-
 3 files changed, 25 insertions(+), 47 deletions(-)

diff --git a/crates/ruff_linter/resources/test/fixtures/pylint/invalid_envvar_default.py b/crates/ruff_linter/resources/test/fixtures/pylint/invalid_envvar_default.py
index 9aa486678e48f..ed1065ce0e233 100644
--- a/crates/ruff_linter/resources/test/fixtures/pylint/invalid_envvar_default.py
+++ b/crates/ruff_linter/resources/test/fixtures/pylint/invalid_envvar_default.py
@@ -6,8 +6,9 @@
 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" + "BAR")
 os.getenv("AA", "GOOD" + 1)
-os.getenv("AA", "GOOD %s" % "BAD")
+os.getenv("AA", "GOOD %s" % "BAR")
 os.getenv("B", Z)
-
+os.getenv("AA", "GOOD" if Z else "BAR")
+os.getenv("AA", 1 if Z else "BAR")  # [invalid-envvar-default]
diff --git a/crates/ruff_linter/src/rules/pylint/rules/invalid_envvar_default.rs b/crates/ruff_linter/src/rules/pylint/rules/invalid_envvar_default.rs
index 4c3b083bfc640..82984a8bc3b49 100644
--- a/crates/ruff_linter/src/rules/pylint/rules/invalid_envvar_default.rs
+++ b/crates/ruff_linter/src/rules/pylint/rules/invalid_envvar_default.rs
@@ -1,6 +1,7 @@
 use ruff_diagnostics::{Diagnostic, Violation};
 use ruff_macros::{derive_message_formats, violation};
-use ruff_python_ast::{self as ast, Expr, Operator};
+use ruff_python_ast as ast;
+use ruff_python_semantic::analyze::type_inference::{PythonType, ResolvedPythonType};
 use ruff_text_size::Ranged;
 
 use crate::checkers::ast::Checker;
@@ -40,43 +41,6 @@ impl Violation for InvalidEnvvarDefault {
     }
 }
 
-fn is_valid_default(expr: &Expr) -> bool {
-    // We can't infer the types of these defaults, so assume they're valid.
-    if matches!(
-        expr,
-        Expr::Name(_) | Expr::Attribute(_) | Expr::Subscript(_) | Expr::Call(_)
-    ) {
-        return true;
-    }
-
-    // Allow string concatenation.
-    if let Expr::BinOp(ast::ExprBinOp {
-        left,
-        right,
-        op: Operator::Add,
-        range: _,
-    }) = expr
-    {
-        return is_valid_default(left) && is_valid_default(right);
-    }
-
-    // Allow string formatting.
-    if let Expr::BinOp(ast::ExprBinOp {
-        left,
-        op: Operator::Mod,
-        ..
-    }) = expr
-    {
-        return is_valid_default(left);
-    }
-
-    // Otherwise, the default must be a string or `None`.
-    matches!(
-        expr,
-        Expr::StringLiteral(_) | Expr::NoneLiteral(_) | Expr::FString(_)
-    )
-}
-
 /// PLW1508
 pub(crate) fn invalid_envvar_default(checker: &mut Checker, call: &ast::ExprCall) {
     if checker
@@ -89,10 +53,15 @@ pub(crate) fn invalid_envvar_default(checker: &mut Checker, call: &ast::ExprCall
             return;
         };
 
-        if !is_valid_default(expr) {
-            checker
-                .diagnostics
-                .push(Diagnostic::new(InvalidEnvvarDefault, expr.range()));
+        if matches!(
+            ResolvedPythonType::from(expr),
+            ResolvedPythonType::Unknown
+                | ResolvedPythonType::Atom(PythonType::String | PythonType::None)
+        ) {
+            return;
         }
+        checker
+            .diagnostics
+            .push(Diagnostic::new(InvalidEnvvarDefault, expr.range()));
     }
 }
diff --git a/crates/ruff_linter/src/rules/pylint/snapshots/ruff_linter__rules__pylint__tests__PLW1508_invalid_envvar_default.py.snap b/crates/ruff_linter/src/rules/pylint/snapshots/ruff_linter__rules__pylint__tests__PLW1508_invalid_envvar_default.py.snap
index ff062bb7005c2..9490733158b2a 100644
--- a/crates/ruff_linter/src/rules/pylint/snapshots/ruff_linter__rules__pylint__tests__PLW1508_invalid_envvar_default.py.snap
+++ b/crates/ruff_linter/src/rules/pylint/snapshots/ruff_linter__rules__pylint__tests__PLW1508_invalid_envvar_default.py.snap
@@ -34,11 +34,19 @@ invalid_envvar_default.py:6:25: PLW1508 Invalid type for environment variable de
 invalid_envvar_default.py:10:17: PLW1508 Invalid type for environment variable default; expected `str` or `None`
    |
  8 | os.getenv("AA", f"GOOD")
- 9 | os.getenv("AA", "GOOD" + "BAD")
+ 9 | os.getenv("AA", "GOOD" + "BAR")
 10 | os.getenv("AA", "GOOD" + 1)
    |                 ^^^^^^^^^^ PLW1508
-11 | os.getenv("AA", "GOOD %s" % "BAD")
+11 | os.getenv("AA", "GOOD %s" % "BAR")
 12 | os.getenv("B", Z)
    |
 
+invalid_envvar_default.py:14:17: PLW1508 Invalid type for environment variable default; expected `str` or `None`
+   |
+12 | os.getenv("B", Z)
+13 | os.getenv("AA", "GOOD" if Z else "BAR")
+14 | os.getenv("AA", 1 if Z else "BAR")  # [invalid-envvar-default]
+   |                 ^^^^^^^^^^^^^^^^^ PLW1508
+   |
+