From 064073a629a97706e4c0cc26548cacc598eb245d Mon Sep 17 00:00:00 2001
From: Kyle Edwards <kyedwards@nvidia.com>
Date: Wed, 28 Aug 2024 14:13:00 -0400
Subject: [PATCH] Fix sorting

---
 src/rapids_pre_commit_hooks/pr_builder.py     | 37 ++++++++++++++-----
 .../test_pr_builder.py                        | 17 +++++----
 2 files changed, 38 insertions(+), 16 deletions(-)

diff --git a/src/rapids_pre_commit_hooks/pr_builder.py b/src/rapids_pre_commit_hooks/pr_builder.py
index 4c4d8ad..36c88a2 100644
--- a/src/rapids_pre_commit_hooks/pr_builder.py
+++ b/src/rapids_pre_commit_hooks/pr_builder.py
@@ -163,21 +163,40 @@ def check_jobs(
                     (pr_builder_key.start_mark.index, pr_builder_key.end_mark.index),
                     "place pr-builder job before all other jobs",
                 )
-                assert self.linter.content[pr_builder_value.end_mark.index] == "\n"
+                end = pr_builder_value.end_mark.index
+                previous_newline = self.linter.content[:end].rfind("\n")
+                if all(
+                    c == " " for c in self.linter.content[previous_newline + 1 : end]
+                ):
+                    end = previous_newline
+                previous_newline_pr_builder = self.linter.content[
+                    : pr_builder_key.start_mark.index
+                ].rfind("\n")
+                assert all(
+                    c == " "
+                    for c in self.linter.content[
+                        previous_newline + 1 : pr_builder_key.start_mark.index
+                    ]
+                )
                 w.add_replacement(
                     (
-                        pr_builder_key.start_mark.index,
-                        pr_builder_value.end_mark.index + 1,
+                        previous_newline_pr_builder + 1,
+                        end + 1,
                     ),
                     "",
                 )
+                previous_newline_first = self.linter.content[
+                    : first_key.start_mark.index
+                ].rfind("\n")
+                assert all(
+                    c == " "
+                    for c in self.linter.content[
+                        previous_newline + 1 : first_key.start_mark.index
+                    ]
+                )
                 w.add_replacement(
-                    (first_key.start_mark.index, first_key.start_mark.index),
-                    self.linter.content[
-                        pr_builder_key.start_mark.index : (
-                            pr_builder_value.end_mark.index + 1
-                        )
-                    ],
+                    (previous_newline_first + 1, previous_newline_first + 1),
+                    self.linter.content[previous_newline_pr_builder + 1 : end + 1],
                 )
 
     def check_root(self) -> None:
diff --git a/test/rapids_pre_commit_hooks/test_pr_builder.py b/test/rapids_pre_commit_hooks/test_pr_builder.py
index c67a7a6..4706680 100644
--- a/test/rapids_pre_commit_hooks/test_pr_builder.py
+++ b/test/rapids_pre_commit_hooks/test_pr_builder.py
@@ -131,7 +131,8 @@ def test_check_pr_builder_job(content, warnings):
             dedent(
                 """\
                 jobs:
-                  pr-builder: {}
+                  pr-builder:
+                    if: always()
                   other-job-1: {}
                   other-job-2: {}
                 """
@@ -144,7 +145,8 @@ def test_check_pr_builder_job(content, warnings):
                 """\
                 jobs:
                   other-job-1: {}
-                  pr-builder: {}
+                  pr-builder:
+                    if: always()
                   other-job-2: {}
                 """
             ),
@@ -154,8 +156,8 @@ def test_check_pr_builder_job(content, warnings):
                     (26, 36),
                     "place pr-builder job before all other jobs",
                     replacements=[
-                        Replacement((26, 41), ""),
-                        Replacement((8, 8), "pr-builder: {}\n"),
+                        Replacement((24, 55), ""),
+                        Replacement((6, 6), "  pr-builder:\n    if: always()\n"),
                     ],
                 ),
             ],
@@ -166,7 +168,8 @@ def test_check_pr_builder_job(content, warnings):
                 jobs:
                   other-job-1: {}
                   other-job-2: {}
-                  pr-builder: {}
+                  pr-builder:
+                    if: always()
                 """
             ),
             2,
@@ -175,8 +178,8 @@ def test_check_pr_builder_job(content, warnings):
                     (44, 54),
                     "place pr-builder job before all other jobs",
                     replacements=[
-                        Replacement((44, 59), ""),
-                        Replacement((8, 8), "pr-builder: {}\n"),
+                        Replacement((42, 73), ""),
+                        Replacement((6, 6), "  pr-builder:\n    if: always()\n"),
                     ],
                 ),
             ],