Skip to main content

ruff_python_trivia/
whitespace.rs

1use std::borrow::Cow;
2
3use ruff_source_file::LineRanges;
4use ruff_text_size::{TextRange, TextSize};
5
6/// Expands tabs to the next eight-column tab stop, matching Python's `str.expandtabs`.
7pub fn expand_tabs(source: &str) -> Cow<'_, str> {
8    const TAB_SIZE: usize = 8;
9
10    if !source.contains('\t') {
11        return Cow::Borrowed(source);
12    }
13
14    let mut expanded = String::with_capacity(source.len());
15    let mut column = 0;
16
17    for character in source.chars() {
18        match character {
19            '\t' => {
20                let spaces = tab_offset(column, TAB_SIZE);
21                expanded.extend(std::iter::repeat_n(' ', spaces));
22                column += spaces;
23            }
24            '\r' | '\n' => {
25                expanded.push(character);
26                column = 0;
27            }
28            _ => {
29                expanded.push(character);
30                column += 1;
31            }
32        }
33    }
34
35    Cow::Owned(expanded)
36}
37
38/// Returns the number of columns from `column` to the next tab stop.
39pub const fn tab_offset(column: usize, tab_size: usize) -> usize {
40    tab_size - column % tab_size
41}
42
43/// Returns the number of columns from `column` to the next tab stop using `u32` values.
44pub const fn tab_offset_u32(column: u32, tab_size: u32) -> u32 {
45    tab_size - column % tab_size
46}
47
48/// Extract the leading indentation from a line.
49pub fn indentation_at_offset(offset: TextSize, source: &str) -> Option<&str> {
50    let line_start = source.line_start(offset);
51    let indentation = &source[TextRange::new(line_start, offset)];
52
53    indentation
54        .chars()
55        .all(is_python_whitespace)
56        .then_some(indentation)
57}
58
59/// Return `true` if the node starting the given [`TextSize`] has leading content.
60pub fn has_leading_content(offset: TextSize, source: &str) -> bool {
61    let line_start = source.line_start(offset);
62    let leading = &source[TextRange::new(line_start, offset)];
63    leading.chars().any(|char| !is_python_whitespace(char))
64}
65
66/// Return `true` if the node ending at the given [`TextSize`] has trailing content.
67pub fn has_trailing_content(offset: TextSize, source: &str) -> bool {
68    let line_end = source.line_end(offset);
69    let trailing = &source[TextRange::new(offset, line_end)];
70
71    for char in trailing.chars() {
72        if char == '#' {
73            return false;
74        }
75        if !is_python_whitespace(char) {
76            return true;
77        }
78    }
79    false
80}
81
82/// Returns `true` for [whitespace](https://docs.python.org/3/reference/lexical_analysis.html#whitespace-between-tokens)
83/// characters.
84pub const fn is_python_whitespace(c: char) -> bool {
85    matches!(
86        c,
87        // Space, tab, or form-feed
88        ' ' | '\t' | '\x0C'
89    )
90}
91
92/// Extract the leading indentation from a line.
93pub fn leading_indentation(line: &str) -> &str {
94    line.find(|char: char| !is_python_whitespace(char))
95        .map_or(line, |index| &line[..index])
96}
97
98pub trait PythonWhitespace {
99    /// Like `str::trim()`, but only removes whitespace characters that Python considers
100    /// to be [whitespace](https://docs.python.org/3/reference/lexical_analysis.html#whitespace-between-tokens).
101    fn trim_whitespace(&self) -> &Self;
102
103    /// Like `str::trim_start()`, but only removes whitespace characters that Python considers
104    /// to be [whitespace](https://docs.python.org/3/reference/lexical_analysis.html#whitespace-between-tokens).
105    fn trim_whitespace_start(&self) -> &Self;
106
107    /// Like `str::trim_end()`, but only removes whitespace characters that Python considers
108    /// to be [whitespace](https://docs.python.org/3/reference/lexical_analysis.html#whitespace-between-tokens).
109    fn trim_whitespace_end(&self) -> &Self;
110}
111
112impl PythonWhitespace for str {
113    fn trim_whitespace(&self) -> &Self {
114        self.trim_matches(is_python_whitespace)
115    }
116
117    fn trim_whitespace_start(&self) -> &Self {
118        self.trim_start_matches(is_python_whitespace)
119    }
120
121    fn trim_whitespace_end(&self) -> &Self {
122        self.trim_end_matches(is_python_whitespace)
123    }
124}
125
126#[cfg(test)]
127mod tests {
128    use std::borrow::Cow;
129
130    use super::{expand_tabs, tab_offset, tab_offset_u32};
131
132    #[test]
133    fn tab_expansion_borrows_unchanged_text() {
134        assert!(matches!(expand_tabs("unchanged"), Cow::Borrowed(_)));
135    }
136
137    #[test]
138    fn tab_expansion_allocates_changed_text() {
139        let expanded = expand_tabs("  \tvalue");
140
141        assert!(matches!(&expanded, Cow::Owned(_)));
142        assert_eq!(expanded, "        value");
143    }
144
145    #[test]
146    fn tab_offset_advances_to_next_stop() {
147        assert_eq!(tab_offset(0, 8), 8);
148        assert_eq!(tab_offset(2, 8), 6);
149        assert_eq!(tab_offset(8, 8), 8);
150    }
151
152    #[test]
153    fn u32_tab_offset_advances_to_next_stop() {
154        assert_eq!(tab_offset_u32(0, 8), 8);
155        assert_eq!(tab_offset_u32(2, 8), 6);
156        assert_eq!(tab_offset_u32(8, 8), 8);
157    }
158}