From a457967e8f6f6ddf6907aba010fb5cbdaad0aa64 Mon Sep 17 00:00:00 2001
From: Amethyst Reese <amethyst@n7.gg>
Date: Tue, 13 Jan 2026 11:16:38 -0800
Subject: [PATCH] Add `split_at` method to `Tokens`

---
 crates/ruff_python_ast/src/token/tokens.rs | 65 ++++++++++++++++++++++
 1 file changed, 65 insertions(+)

diff --git a/crates/ruff_python_ast/src/token/tokens.rs b/crates/ruff_python_ast/src/token/tokens.rs
index 7ca096d9b0..26f4d88649 100644
--- a/crates/ruff_python_ast/src/token/tokens.rs
+++ b/crates/ruff_python_ast/src/token/tokens.rs
@@ -185,6 +185,31 @@ impl Tokens {
 
         after
     }
+
+    /// Returns a pair of token slices from both before and after the given [`TextSize`] offset.
+    ///
+    /// If the given offset is between two tokens, the "before" slice will end just before the
+    /// following token. In other words, if the offset is between the end of previous token and
+    /// start of next token, the "before" slice will end just before the next token. The "after"
+    /// slice will contain the rest of the tokens.
+    ///
+    /// # Panics
+    ///
+    /// If the given offset is inside a token range at any point
+    /// other than the start of the range.
+    pub fn split_at(&self, offset: TextSize) -> (&[Token], &[Token]) {
+        let partition_point = self.partition_point(|token| token.start() < offset);
+        let before = &self.raw[..partition_point];
+        let after = &self.raw[partition_point..];
+
+        if let Some(last) = before.last() {
+            assert!(
+                offset >= last.end(),
+                "Offset {offset:?} is inside token `{last:?}`"
+            );
+        }
+        (before, after)
+    }
 }
 
 impl<'a> IntoIterator for &'a Tokens {
@@ -513,4 +538,44 @@ mod tests {
         let tokens = new_tokens(TEST_CASE_WITH_GAP.into_iter());
         tokens.in_range(TextRange::new(0.into(), 6.into()));
     }
+
+    #[test]
+    fn tokens_split_at_first_token_start() {
+        let tokens = new_tokens(TEST_CASE_WITH_GAP.into_iter());
+        let (before, after) = tokens.split_at(TextSize::new(0));
+        assert_eq!(before.len(), 0);
+        assert_eq!(after.len(), 10);
+    }
+
+    #[test]
+    fn tokens_split_at_last_token_end() {
+        let tokens = new_tokens(TEST_CASE_WITH_GAP.into_iter());
+        let (before, after) = tokens.split_at(TextSize::new(33));
+        assert_eq!(before.len(), 10);
+        assert_eq!(after.len(), 0);
+    }
+
+    #[test]
+    fn tokens_split_at_inside_gap() {
+        let tokens = new_tokens(TEST_CASE_WITH_GAP.into_iter());
+        let (before, after) = tokens.split_at(TextSize::new(13));
+        assert_eq!(before.len(), 6);
+        assert_eq!(after.len(), 4);
+    }
+
+    #[test]
+    #[should_panic(expected = "Offset 18 is inside token `Comment 15..24`")]
+    fn tokens_split_at_inside_token() {
+        let tokens = new_tokens(TEST_CASE_WITH_GAP.into_iter());
+        tokens.split_at(TextSize::new(18));
+    }
+
+    #[test]
+    fn tokens_split_at_matches_before_and_after() {
+        let offset = TextSize::new(15);
+        let tokens = new_tokens(TEST_CASE_WITH_GAP.into_iter());
+        let (before, after) = tokens.split_at(offset);
+        assert_eq!(before, tokens.before(offset));
+        assert_eq!(after, tokens.after(offset));
+    }
 }