Support IfExp with dual string arms in invalid-envvar-value

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

@@ -10,3 +10,4 @@
 os.getenv("AA", "GOOD" + 1)
 os.getenv("AA", "GOOD %s" % "BAD")
 os.getenv("B", Z)
+

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

@@ -10,6 +10,8 @@
 os.getenv(key=f"foo", default="bar")
 os.getenv(key="foo" + "bar", default=1)
 os.getenv(key=1 + "bar", default=1)  # [invalid-envvar-value]
+os.getenv("PATH_TEST" if using_clear_path else "PATH_ORIG")
+os.getenv(1 if using_clear_path else "PATH_ORIG")
 
 AA = "aa"
 os.getenv(AA)

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

@@ -1,6 +1,7 @@
 use ruff_diagnostics::{Diagnostic, Violation};
 use ruff_macros::{derive_message_formats, violation};
-use ruff_python_ast::{self as ast, Constant, Expr, Operator, Ranged};
+use ruff_python_ast::{self as ast, Ranged};
+use ruff_python_semantic::analyze::type_inference::PythonType;
 
 use crate::checkers::ast::Checker;
 
@@ -32,46 +33,6 @@ 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,
-        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_key(left) && is_valid_key(right);
-    }
-
-    // Allow string formatting.
-    if let Expr::BinOp(ast::ExprBinOp {
-        left,
-        op: Operator::Mod,
-        ..
-    }) = expr
-    {
-        return is_valid_key(left);
-    }
-
-    // Otherwise, the default must be a string.
-    matches!(
-        expr,
-        Expr::Constant(ast::ExprConstant {
-            value: Constant::Str { .. },
-            ..
-        }) | Expr::FString(_)
-    )
-}
-
 /// PLE1507
 pub(crate) fn invalid_envvar_value(checker: &mut Checker, call: &ast::ExprCall) {
     if checker
@@ -84,10 +45,15 @@ pub(crate) fn invalid_envvar_value(checker: &mut Checker, call: &ast::ExprCall)
             return;
         };
 
-        if !is_valid_key(expr) {
-            checker
-                .diagnostics
-                .push(Diagnostic::new(InvalidEnvvarValue, expr.range()));
+        if matches!(
+            PythonType::from(expr),
+            PythonType::String | PythonType::Unknown
+        ) {
+            return;
         }
+
+        checker
+            .diagnostics
+            .push(Diagnostic::new(InvalidEnvvarValue, expr.range()));
     }
 }

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

@@ -41,7 +41,6 @@ pub(crate) fn invalid_str_return(checker: &mut Checker, name: &str, body: &[Stmt
 
     for stmt in returns {
         if let Some(value) = stmt.value.as_deref() {
-            // Disallow other, non-
             if !matches!(
                 PythonType::from(value),
                 PythonType::String | PythonType::Unknown

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

@@ -31,14 +31,4 @@ invalid_envvar_value.py:8:11: PLE1507 Invalid type for initial `os.getenv` argum
 10 | os.getenv(key=f"foo", default="bar")
    |
 
-invalid_envvar_value.py:12:15: PLE1507 Invalid type for initial `os.getenv` argument; expected `str`
-   |
-10 | os.getenv(key=f"foo", default="bar")
-11 | os.getenv(key="foo" + "bar", default=1)
-12 | os.getenv(key=1 + "bar", default=1)  # [invalid-envvar-value]
-   |               ^^^^^^^^^ PLE1507
-13 | 
-14 | AA = "aa"
-   |
-
 

crates/ruff_python_semantic/src/analyze/type_inference.rs

@@ -1,15 +1,15 @@
 //! Analysis rules to perform basic type inference on individual expressions.
 
 use ruff_python_ast as ast;
-use ruff_python_ast::{Constant, Expr};
+use ruff_python_ast::{Constant, Expr, Operator};
 
 /// An extremely simple type inference system for individual expressions.
 ///
 /// This system can only represent and infer the types of simple data types
 /// such as strings, integers, floats, and containers. It cannot infer the
 /// types of variables or expressions that are not statically known from
 /// individual AST nodes alone.
-#[derive(Debug, Copy, Clone)]
+#[derive(Debug, Copy, Clone, PartialEq, Eq, is_macro::Is)]
 pub enum PythonType {
     /// A string literal, such as `"hello"`.
     String,
@@ -55,27 +55,43 @@ impl From<&Expr> for PythonType {
             Expr::Tuple(_) => PythonType::Tuple,
             Expr::GeneratorExp(_) => PythonType::Generator,
             Expr::FString(_) => PythonType::String,
-            Expr::BinOp(ast::ExprBinOp { left, op, .. }) => {
-                // Ex) "a" % "b"
-                if op.is_mod() {
-                    if matches!(
-                        left.as_ref(),
-                        Expr::Constant(ast::ExprConstant {
-                            value: Constant::Str(..),
-                            ..
-                        })
-                    ) {
-                        return PythonType::String;
-                    }
-                    if matches!(
-                        left.as_ref(),
-                        Expr::Constant(ast::ExprConstant {
-                            value: Constant::Bytes(..),
-                            ..
-                        })
-                    ) {
-                        return PythonType::Bytes;
+            Expr::IfExp(ast::ExprIfExp { body, orelse, .. }) => {
+                let body = PythonType::from(body.as_ref());
+                let orelse = PythonType::from(orelse.as_ref());
+                // TODO(charlie): If we have two known types, we should return a union. As-is,
+                // callers that ignore the `Unknown` type will allow invalid expressions (e.g.,
+                // if you're testing for strings, you may accept `String` or `Unknown`, and you'd
+                // now accept, e.g., `1 if True else "a"`, which resolves to `Unknown`).
+                if body == orelse {
+                    body
+                } else {
+                    PythonType::Unknown
+                }
+            }
+            Expr::BinOp(ast::ExprBinOp {
+                left, op, right, ..
+            }) => {
+                match op {
+                    // Ex) "a" + "b"
+                    Operator::Add => {
+                        match (
+                            PythonType::from(left.as_ref()),
+                            PythonType::from(right.as_ref()),
+                        ) {
+                            (PythonType::String, PythonType::String) => return PythonType::String,
+                            (PythonType::Bytes, PythonType::Bytes) => return PythonType::Bytes,
+                            // TODO(charlie): If we have two known types, they may be incompatible.
+                            // Return an error (e.g., for `1 + "a"`).
+                            _ => {}
+                        }
                     }
+                    // Ex) "a" % "b"
+                    Operator::Mod => match PythonType::from(left.as_ref()) {
+                        PythonType::String => return PythonType::String,
+                        PythonType::Bytes => return PythonType::Bytes,
+                        _ => {}
+                    },
+                    _ => {}
                 }
                 PythonType::Unknown
             }