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
  }