Autofix PIE810 rule violations (#3411)

crates/ruff/resources/test/fixtures/flake8_pie/PIE810.py

@@ -6,6 +6,8 @@ obj.endswith("foo") or obj.endswith("bar")
 obj.startswith(foo) or obj.startswith(bar)
 # error
 obj.startswith(foo) or obj.startswith("foo")
+# error
+obj.endswith(foo) or obj.startswith(foo) or obj.startswith("foo")
 
 # ok
 obj.startswith(("foo",  "bar"))

crates/ruff/src/checkers/ast/mod.rs

@@ -3435,7 +3435,7 @@ where
                     pylint::rules::merge_isinstance(self, expr, op, values);
                 }
                 if self.settings.rules.enabled(&Rule::SingleStartsEndsWith) {
-                    flake8_pie::rules::single_starts_ends_with(self, values, op);
+                    flake8_pie::rules::single_starts_ends_with(self, expr);
                 }
                 if self.settings.rules.enabled(&Rule::DuplicateIsinstanceCall) {
                     flake8_simplify::rules::duplicate_isinstance_call(self, expr);

crates/ruff/src/rules/flake8_pie/rules.rs

@@ -1,12 +1,18 @@
+use itertools::Either::{Left, Right};
+use std::collections::BTreeMap;
+use std::iter;
+
 use log::error;
 use rustc_hash::FxHashSet;
-use rustpython_parser::ast::{Boolop, Constant, Expr, ExprKind, Keyword, Stmt, StmtKind};
+use rustpython_parser::ast::{
+    Boolop, Constant, Expr, ExprContext, ExprKind, Keyword, Stmt, StmtKind,
+};
 
 use ruff_diagnostics::{AlwaysAutofixableViolation, Violation};
 use ruff_diagnostics::{Diagnostic, Fix};
 use ruff_macros::{derive_message_formats, violation};
 use ruff_python_ast::comparable::ComparableExpr;
-use ruff_python_ast::helpers::{match_trailing_comment, unparse_expr};
+use ruff_python_ast::helpers::{create_expr, match_trailing_comment, unparse_expr};
 use ruff_python_ast::types::{Range, RefEquality};
 use ruff_python_stdlib::identifiers::is_identifier;
 use ruff_python_stdlib::keyword::KWLIST;
@@ -120,12 +126,17 @@ pub struct SingleStartsEndsWith {
     pub attr: String,
 }
 
-impl Violation for SingleStartsEndsWith {
+impl AlwaysAutofixableViolation for SingleStartsEndsWith {
     #[derive_message_formats]
     fn message(&self) -> String {
         let SingleStartsEndsWith { attr } = self;
         format!("Call `{attr}` once with a `tuple`")
     }
+
+    fn autofix_title(&self) -> String {
+        let SingleStartsEndsWith { attr } = self;
+        format!("Merge into a single `{attr}` call")
+    }
 }
 
 #[violation]
@@ -392,39 +403,116 @@ pub fn no_unnecessary_dict_kwargs(checker: &mut Checker, expr: &Expr, kwargs: &[
 }
 
 /// PIE810
-pub fn single_starts_ends_with(checker: &mut Checker, values: &[Expr], node: &Boolop) {
+pub fn single_starts_ends_with(checker: &mut Checker, expr: &Expr) {
-    if *node != Boolop::Or {
+    let ExprKind::BoolOp { op: Boolop::Or, values } = &expr.node else {
         return;
-    }
+    };
 
-    // Given `foo.startswith`, insert ("foo", "startswith") into the set.
+    let mut duplicates = BTreeMap::new();
-    let mut seen = FxHashSet::default();
+    for (index, call) in values.iter().enumerate() {
-    for expr in values {
+        let ExprKind::Call {
-        if let ExprKind::Call {
             func,
             args,
             keywords,
             ..
-        } = &expr.node
+        } = &call.node else {
-        {
+            continue
-            if !(args.len() == 1 && keywords.is_empty()) {
+        };
-                continue;
+
-            }
+        if !(args.len() == 1 && keywords.is_empty()) {
-            if let ExprKind::Attribute { value, attr, .. } = &func.node {
+            continue;
-                if attr != "startswith" && attr != "endswith" {
+        }
-                    continue;
+
-                }
+        let ExprKind::Attribute { value, attr, .. } = &func.node else {
-                if let ExprKind::Name { id, .. } = &value.node {
+            continue
-                    if !seen.insert((id, attr)) {
+        };
-                        checker.diagnostics.push(Diagnostic::new(
+
-                            SingleStartsEndsWith {
+        if attr != "startswith" && attr != "endswith" {
-                                attr: attr.to_string(),
+            continue;
-                            },
+        }
-                            Range::from(value),
+
-                        ));
+        let ExprKind::Name { id: arg_name, .. } = &value.node else {
-                    }
+            continue
-                }
+        };
+
+        duplicates
+            .entry((attr.as_str(), arg_name.as_str()))
+            .or_insert_with(Vec::new)
+            .push(index);
+    }
+
+    // Generate a `Diagnostic` for each duplicate.
+    for ((attr_name, arg_name), indices) in duplicates {
+        if indices.len() > 1 {
+            let mut diagnostic = Diagnostic::new(
+                SingleStartsEndsWith {
+                    attr: attr_name.to_string(),
+                },
+                Range::from(expr),
+            );
+            if checker.patch(diagnostic.kind.rule()) {
+                let words: Vec<&Expr> = indices
+                    .iter()
+                    .map(|index| &values[*index])
+                    .map(|expr| {
+                        let ExprKind::Call { func: _, args, keywords: _} = &expr.node else {
+                            unreachable!("{}", format!("Indices should only contain `{attr_name}` calls"))
+                        };
+                        args.get(0)
+                            .unwrap_or_else(|| panic!("`{attr_name}` should have one argument"))
+                    })
+                    .collect();
+
+                let call = create_expr(ExprKind::Call {
+                    func: Box::new(create_expr(ExprKind::Attribute {
+                        value: Box::new(create_expr(ExprKind::Name {
+                            id: arg_name.to_string(),
+                            ctx: ExprContext::Load,
+                        })),
+                        attr: attr_name.to_string(),
+                        ctx: ExprContext::Load,
+                    })),
+                    args: vec![create_expr(ExprKind::Tuple {
+                        elts: words
+                            .iter()
+                            .flat_map(|value| {
+                                if let ExprKind::Tuple { elts, .. } = &value.node {
+                                    Left(elts.iter())
+                                } else {
+                                    Right(iter::once(*value))
+                                }
+                            })
+                            .map(Clone::clone)
+                            .collect(),
+                        ctx: ExprContext::Load,
+                    })],
+                    keywords: vec![],
+                });
+
+                // Generate the combined `BoolOp`.
+                let mut call = Some(call);
+                let bool_op = create_expr(ExprKind::BoolOp {
+                    op: Boolop::Or,
+                    values: values
+                        .iter()
+                        .enumerate()
+                        .filter_map(|(index, elt)| {
+                            if indices.contains(&index) {
+                                std::mem::take(&mut call)
+                            } else {
+                                Some(elt.clone())
+                            }
+                        })
+                        .collect(),
+                });
+
+                diagnostic.amend(Fix::replacement(
+                    unparse_expr(&bool_op, checker.stylist),
+                    expr.location,
+                    expr.end_location.unwrap(),
+                ));
             }
+            checker.diagnostics.push(diagnostic);
         }
     }
 }

crates/ruff/src/rules/flake8_pie/snapshots/ruff__rules__flake8_pie__tests__PIE810_PIE810.py.snap

@@ -5,53 +5,101 @@ expression: diagnostics
 - kind:
     name: SingleStartsEndsWith
     body: "Call `startswith` once with a `tuple`"
-    suggestion: ~
+    suggestion: "Merge into a single `startswith` call"
-    fixable: false
+    fixable: true
   location:
     row: 2
-    column: 25
+    column: 0
   end_location:
     row: 2
-    column: 28
+    column: 46
-  fix: ~
+  fix:
+    content: "obj.startswith((\"foo\", \"bar\"))"
+    location:
+      row: 2
+      column: 0
+    end_location:
+      row: 2
+      column: 46
   parent: ~
 - kind:
     name: SingleStartsEndsWith
     body: "Call `endswith` once with a `tuple`"
-    suggestion: ~
+    suggestion: "Merge into a single `endswith` call"
-    fixable: false
+    fixable: true
   location:
     row: 4
-    column: 23
+    column: 0
   end_location:
     row: 4
-    column: 26
+    column: 42
-  fix: ~
+  fix:
+    content: "obj.endswith((\"foo\", \"bar\"))"
+    location:
+      row: 4
+      column: 0
+    end_location:
+      row: 4
+      column: 42
   parent: ~
 - kind:
     name: SingleStartsEndsWith
     body: "Call `startswith` once with a `tuple`"
-    suggestion: ~
+    suggestion: "Merge into a single `startswith` call"
-    fixable: false
+    fixable: true
   location:
     row: 6
-    column: 23
+    column: 0
   end_location:
     row: 6
-    column: 26
+    column: 42
-  fix: ~
+  fix:
+    content: "obj.startswith((foo, bar))"
+    location:
+      row: 6
+      column: 0
+    end_location:
+      row: 6
+      column: 42
   parent: ~
 - kind:
     name: SingleStartsEndsWith
     body: "Call `startswith` once with a `tuple`"
-    suggestion: ~
+    suggestion: "Merge into a single `startswith` call"
-    fixable: false
+    fixable: true
   location:
     row: 8
-    column: 23
+    column: 0
   end_location:
     row: 8
-    column: 26
+    column: 44
-  fix: ~
+  fix:
+    content: "obj.startswith((foo, \"foo\"))"
+    location:
+      row: 8
+      column: 0
+    end_location:
+      row: 8
+      column: 44
+  parent: ~
+- kind:
+    name: SingleStartsEndsWith
+    body: "Call `startswith` once with a `tuple`"
+    suggestion: "Merge into a single `startswith` call"
+    fixable: true
+  location:
+    row: 10
+    column: 0
+  end_location:
+    row: 10
+    column: 65
+  fix:
+    content: "obj.endswith(foo) or obj.startswith((foo, \"foo\"))"
+    location:
+      row: 10
+      column: 0
+    end_location:
+      row: 10
+      column: 65
   parent: ~