[pylint]: call swap_with_temporary_variable from definition analysis step

Author: Bnyro

crates/ruff_linter/src/checkers/ast/analyze/definitions.rs

@@ -17,6 +17,33 @@ use crate::{docstrings, warn_user};
 /// it is expected that all [`Definition`] nodes have been visited by the time, and that this
 /// method will not recurse into any other nodes.
 pub(crate) fn definitions(checker: &mut Checker) {
+    // Compute visibility of all definitions.
+    let exports: Option<Vec<DunderAllName>> = {
+        checker
+            .semantic
+            .global_scope()
+            .get_all("__all__")
+            .map(|binding_id| &checker.semantic.bindings[binding_id])
+            .filter_map(|binding| match &binding.kind {
+                BindingKind::Export(Export { names }) => Some(names.iter().copied()),
+                _ => None,
+            })
+            .fold(None, |acc, names| {
+                Some(acc.into_iter().flatten().chain(names).collect())
+            })
+    };
+
+    let definitions = std::mem::take(&mut checker.semantic.definitions);
+    for ContextualizedDefinition {
+        definition,
+        visibility: _,
+    } in definitions.resolve(exports.as_deref()).iter()
+    {
+        if checker.is_rule_enabled(Rule::SwapWithTemporaryVariable) {
+            pylint::rules::swap_with_temporary_variable(checker, definition);
+        }
+    }
+
     let enforce_annotations = checker.any_rule_enabled(&[
         Rule::AnyType,
         Rule::MissingReturnTypeClassMethod,
@@ -100,24 +127,8 @@ pub(crate) fn definitions(checker: &mut Checker) {
         return;
     }
 
-    // Compute visibility of all definitions.
-    let exports: Option<Vec<DunderAllName>> = {
-        checker
-            .semantic
-            .global_scope()
-            .get_all("__all__")
-            .map(|binding_id| &checker.semantic.bindings[binding_id])
-            .filter_map(|binding| match &binding.kind {
-                BindingKind::Export(Export { names }) => Some(names.iter().copied()),
-                _ => None,
-            })
-            .fold(None, |acc, names| {
-                Some(acc.into_iter().flatten().chain(names).collect())
-            })
-    };
-
-    let definitions = std::mem::take(&mut checker.semantic.definitions);
     let mut overloaded_name: Option<&str> = None;
+    let definitions = std::mem::take(&mut checker.semantic.definitions);
     for ContextualizedDefinition {
         definition,
         visibility,

crates/ruff_linter/src/checkers/ast/analyze/statement.rs

@@ -1476,9 +1476,6 @@ pub(crate) fn statement(stmt: &Stmt, checker: &mut Checker) {
             if checker.is_rule_enabled(Rule::NonPEP695TypeAlias) {
                 pyupgrade::rules::non_pep695_type_alias_type(checker, assign);
             }
-            if checker.is_rule_enabled(Rule::SwapWithTemporaryVariable) {
-                pylint::rules::swap_with_temporary_variable(checker, stmt);
-            }
         }
         Stmt::AnnAssign(
             assign_stmt @ ast::StmtAnnAssign {
@@ -1528,9 +1525,6 @@ pub(crate) fn statement(stmt: &Stmt, checker: &mut Checker) {
             if checker.is_rule_enabled(Rule::UnsortedDunderAll) {
                 ruff::rules::sort_dunder_all_ann_assign(checker, assign_stmt);
             }
-            if checker.is_rule_enabled(Rule::SwapWithTemporaryVariable) {
-                pylint::rules::swap_with_temporary_variable(checker, stmt);
-            }
             if checker.source_type.is_stub() {
                 if let Some(value) = value {
                     if checker.is_rule_enabled(Rule::AssignmentDefaultInStub) {

crates/ruff_linter/src/rules/pylint/rules/swap_with_temporary_variable.rs

@@ -1,9 +1,8 @@
 use itertools::Itertools;
 use ruff_diagnostics::{Applicability, Edit, Fix};
+use ruff_python_ast::Stmt;
 use ruff_python_ast::name::Name;
-use ruff_python_ast::traversal::EnclosingSuite;
-use ruff_python_ast::{Stmt, traversal};
-use ruff_python_semantic::SemanticModel;
+use ruff_python_semantic::Definition;
 use ruff_text_size::{Ranged, TextRange, TextSize};
 
 use ruff_macros::{ViolationMetadata, derive_message_formats};
@@ -42,7 +41,7 @@ use crate::checkers::ast::Checker;
 /// exception case, applying the quick fix would remove comments between the
 /// assignment statements.
 #[derive(ViolationMetadata)]
-#[violation_metadata(preview_since = "0.14.11")]
+#[violation_metadata(preview_since = "0.15.1")]
 pub(crate) struct SwapWithTemporaryVariable<'a> {
     first_var: &'a Name,
     second_var: &'a Name,
@@ -79,13 +78,14 @@ impl Violation for SwapWithTemporaryVariable<'_> {
     }
 }
 
-pub(crate) fn swap_with_temporary_variable(checker: &Checker, assignment: &Stmt) {
-    let Some(consecutive_assignments) =
-        match_consecutive_assignments(assignment, checker.semantic())
-    else {
-        return;
-    };
-    for (stmt_a, stmt_b, stmt_c) in consecutive_assignments.tuple_windows() {
+pub(crate) fn swap_with_temporary_variable(checker: &Checker, definition: &Definition) {
+    let consecutive_assignments = match_consecutive_assignments(definition);
+
+    for stmt_sequence in consecutive_assignments.tuple_windows() {
+        let (Some(stmt_a), Some(stmt_b), Some(stmt_c)) = stmt_sequence else {
+            continue;
+        };
+
         // Detect patterns like:
         // temp = x
         // x = y
@@ -134,41 +134,14 @@ pub(crate) fn swap_with_temporary_variable(checker: &Checker, assignment: &Stmt)
 ///
 /// Also see the `repeated_append` rule for a similar use case.
 fn match_consecutive_assignments<'a>(
-    stmt: &'a Stmt,
-    semantic: &'a SemanticModel,
-) -> Option<impl Iterator<Item = VarToVarAssignment<'a>>> {
-    let root_assignment = VarToVarAssignment::from_stmt(stmt)?;
-
-    // In order to match consecutive statements, we need to go to the tree ancestor of the
-    // given statement, find its position there, and match all 'appends' from there.
-    let suite = if semantic.at_top_level() {
-        // If the statement is at the top level, we should go to the parent module.
-        // Module is available in the definitions list.
-        EnclosingSuite::new(semantic.definitions.python_ast()?, stmt.into())?
-    } else {
-        // Otherwise, go to the parent, and take its body as a sequence of siblings.
-        semantic
-            .current_statement_parent()
-            .and_then(|parent| traversal::suite(stmt, parent))?
+    definition: &'a Definition,
+) -> impl Iterator<Item = Option<VarToVarAssignment<'a>>> {
+    let suite = match definition {
+        Definition::Module(module) => module.python_ast,
+        Definition::Member(member) => member.body(),
     };
 
-    // We shouldn't repeat the same work for many 'assignments' that go in a row. Let's check
-    // that this statement is at the beginning of such a group.
-    if suite
-        .previous_sibling()
-        .is_some_and(|previous_stmt| VarToVarAssignment::from_stmt(previous_stmt).is_some())
-    {
-        return None;
-    }
-
-    Some(
-        std::iter::once(root_assignment).chain(
-            suite
-                .next_siblings()
-                .iter()
-                .map_while(VarToVarAssignment::from_stmt),
-        ),
-    )
+    suite.iter().map(VarToVarAssignment::from_stmt)
 }
 
 #[derive(Eq, PartialEq, Debug, Clone)]