Repository: astral-sh/ruff
File: crates/ruff_python_ast/src/token/tokens.rs (Rust)

Code before:
    ///
    /// If the given offset is inside a token range at any point
    /// other than the start of the range.
    pub fn after(&self, offset: TextSize) -> &[Token] {
        let partition_point = self.partition_point(|token| token.end() <= offset);
        let after = &self[partition_point..];

        if let Some(first) = after.first() {
            // valid. If it's greater than the end offset, then it's in the gap between
            // the tokens which is valid as well.
            assert!(
                offset <= first.start(),
                "Offset {offset:?} is inside token `{first:?}`",
            );
        }

        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 {
    type Item = &'a Token;
    type IntoIter = std::slice::Iter<'a, Token>;

    fn into_iter(self) -> Self::IntoIter {
        self.iter()

Review comment (suggestion):
I think it's worth pointing out that `after` is different than calling `.after` (at least I think so).

**`.after`**

Finds the first token that ends after `offset`

```
           
| 1 |  | 2 |     | 3|
           - offset
                 ^^^^ .after
^^^^^ split_at before
       ^^^^^^^^^^^^^^ split_at after        
```

This is because `.after` finds the first token that ends after `offset` whereas `split_at` finds the first token that starts at or before `offset`

We could change `split_at` to skip over tokens that fall directly on `offset` (e.g. a `Dedent` token with zero width). For now, I think it's fine to call this difference out in the comment

Code after:
    ///
    /// If the given offset is inside a token range at any point
    /// other than the start of the range.
    pub fn after(&self, offset: TextSize) -> &[Token] {
        let partition_point = self.partition_point(|token| token.end() <= offset);
        let after = &self[partition_point..];

        if let Some(first) = after.first() {
            // valid. If it's greater than the end offset, then it's in the gap between
            // the tokens which is valid as well.
            assert!(
                offset <= first.start(),
                "Offset {offset:?} is inside token `{first:?}`",
            );
        }

        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.
    ///
    /// Note that the contents of the "after" slice may differ from the results of calling `after()`
    /// directly, particularly when the given offset occurs on zero-width tokens like `Dedent`.
    ///
    /// # 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, after) = &self.raw.split_at(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 {
    type Item = &'a Token;
    type IntoIter = std::slice::Iter<'a, Token>;