Reduce memory usage of Docstring struct (#16183)

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

@@ -1,10 +1,8 @@
-use ruff_python_ast::str::raw_contents_range;
 use ruff_python_semantic::all::DunderAllName;
 use ruff_python_semantic::{
     BindingKind, ContextualizedDefinition, Definition, Export, Member, MemberKind,
 };
-use ruff_source_file::LineRanges;
-use ruff_text_size::{Ranged, TextRange};
+use ruff_text_size::Ranged;
 
 use crate::checkers::ast::Checker;
 use crate::codes::Rule;
@@ -184,14 +182,9 @@ pub(crate) fn definitions(checker: &mut Checker) {
                 continue;
             };
 
-            let contents = checker.locator().slice(string_literal);
-
-            let indentation = checker.locator().slice(TextRange::new(
-                checker.locator.line_start(string_literal.start()),
-                string_literal.start(),
-            ));
-
-            if string_literal.value.is_implicit_concatenated() {
+            // If the `ExprStringLiteral` has multiple parts, it is implicitly concatenated.
+            // We don't support recognising such strings as docstrings in our model currently.
+            let [sole_string_part] = string_literal.value.as_slice() else {
                 #[allow(deprecated)]
                 let location = checker
                     .locator
@@ -203,16 +196,12 @@ pub(crate) fn definitions(checker: &mut Checker) {
                     location.column
                 );
                 continue;
-            }
+            };
 
-            // SAFETY: Safe for docstrings that pass `should_ignore_docstring`.
-            let body_range = raw_contents_range(contents).unwrap();
             let docstring = Docstring {
                 definition,
-                expr: string_literal,
-                contents,
-                body_range,
-                indentation,
+                expr: sole_string_part,
+                source: checker.source(),
             };
 
             if !pydocstyle::rules::not_empty(checker, &docstring) {

crates/ruff_linter/src/docstrings/mod.rs

@@ -1,9 +1,10 @@
 use std::fmt::{Debug, Formatter};
 use std::ops::Deref;
 
-use ruff_python_ast::ExprStringLiteral;
+use ruff_python_ast::{self as ast, StringFlags};
 use ruff_python_semantic::Definition;
-use ruff_text_size::{Ranged, TextRange};
+use ruff_source_file::LineRanges;
+use ruff_text_size::{Ranged, TextRange, TextSize};
 
 pub(crate) mod extraction;
 pub(crate) mod google;
@@ -15,26 +16,71 @@ pub(crate) mod styles;
 pub(crate) struct Docstring<'a> {
     pub(crate) definition: &'a Definition<'a>,
     /// The literal AST node representing the docstring.
-    pub(crate) expr: &'a ExprStringLiteral,
-    /// The content of the docstring, including the leading and trailing quotes.
-    pub(crate) contents: &'a str,
-    /// The range of the docstring body (without the quotes). The range is relative to [`Self::contents`].
-    pub(crate) body_range: TextRange,
-    pub(crate) indentation: &'a str,
+    pub(crate) expr: &'a ast::StringLiteral,
+    /// The source file the docstring was defined in.
+    pub(crate) source: &'a str,
 }
 
 impl<'a> Docstring<'a> {
+    fn flags(&self) -> ast::StringLiteralFlags {
+        self.expr.flags
+    }
+
+    /// The contents of the docstring, including the opening and closing quotes.
+    pub(crate) fn contents(&self) -> &'a str {
+        &self.source[self.range()]
+    }
+
+    /// The contents of the docstring, excluding the opening and closing quotes.
     pub(crate) fn body(&self) -> DocstringBody {
         DocstringBody { docstring: self }
     }
 
-    pub(crate) fn leading_quote(&self) -> &'a str {
-        &self.contents[TextRange::up_to(self.body_range.start())]
+    /// Compute the start position of the docstring's opening line
+    pub(crate) fn line_start(&self) -> TextSize {
+        self.source.line_start(self.start())
     }
 
-    pub(crate) fn triple_quoted(&self) -> bool {
-        let leading_quote = self.leading_quote();
-        leading_quote.ends_with("\"\"\"") || leading_quote.ends_with("'''")
+    /// Return the slice of source code that represents the indentation of the docstring's opening quotes.
+    pub(crate) fn compute_indentation(&self) -> &'a str {
+        &self.source[TextRange::new(self.line_start(), self.start())]
+    }
+
+    pub(crate) fn quote_style(&self) -> ast::str::Quote {
+        self.flags().quote_style()
+    }
+
+    pub(crate) fn is_raw_string(&self) -> bool {
+        self.flags().prefix().is_raw()
+    }
+
+    pub(crate) fn is_u_string(&self) -> bool {
+        self.flags().prefix().is_unicode()
+    }
+
+    pub(crate) fn is_triple_quoted(&self) -> bool {
+        self.flags().is_triple_quoted()
+    }
+
+    /// The docstring's prefixes as they exist in the original source code.
+    pub(crate) fn prefix_str(&self) -> &'a str {
+        // N.B. This will normally be exactly the same as what you might get from
+        // `self.flags().prefix().as_str()`, but doing it this way has a few small advantages.
+        // For example, the casing of the `u` prefix will be preserved if it's a u-string.
+        &self.source[TextRange::new(
+            self.start(),
+            self.start() + self.flags().prefix().text_len(),
+        )]
+    }
+
+    /// The docstring's "opener" (the string's prefix, if any, and its opening quotes).
+    pub(crate) fn opener(&self) -> &'a str {
+        &self.source[TextRange::new(self.start(), self.start() + self.flags().opener_len())]
+    }
+
+    /// The docstring's closing quotes.
+    pub(crate) fn closer(&self) -> &'a str {
+        &self.source[TextRange::new(self.end() - self.flags().closer_len(), self.end())]
     }
 }
 
@@ -51,13 +97,13 @@ pub(crate) struct DocstringBody<'a> {
 
 impl<'a> DocstringBody<'a> {
     pub(crate) fn as_str(self) -> &'a str {
-        &self.docstring.contents[self.docstring.body_range]
+        &self.docstring.source[self.range()]
     }
 }
 
 impl Ranged for DocstringBody<'_> {
     fn range(&self) -> TextRange {
-        self.docstring.body_range + self.docstring.start()
+        self.docstring.expr.content_range()
     }
 }
 

crates/ruff_linter/src/rules/pydocstyle/rules/backslashes.rs

@@ -59,8 +59,7 @@ impl Violation for EscapeSequenceInDocstring {
 
 /// D301
 pub(crate) fn backslashes(checker: &Checker, docstring: &Docstring) {
-    // Docstring is already raw.
-    if docstring.leading_quote().contains(['r', 'R']) {
+    if docstring.is_raw_string() {
         return;
     }
 
@@ -99,10 +98,10 @@ pub(crate) fn backslashes(checker: &Checker, docstring: &Docstring) {
         if !matches!(*escaped_char, '\r' | '\n' | 'u' | 'U' | 'N') {
             let mut diagnostic = Diagnostic::new(EscapeSequenceInDocstring, docstring.range());
 
-            if !docstring.leading_quote().contains(['u', 'U']) {
-                diagnostic.set_fix(Fix::unsafe_edit(Edit::range_replacement(
-                    "r".to_owned() + docstring.contents,
-                    docstring.range(),
+            if !docstring.is_u_string() {
+                diagnostic.set_fix(Fix::unsafe_edit(Edit::insertion(
+                    "r".to_string(),
+                    docstring.start(),
                 )));
             }
 

crates/ruff_linter/src/rules/pydocstyle/rules/blank_after_summary.rs

@@ -69,7 +69,7 @@ impl Violation for MissingBlankLineAfterSummary {
 pub(crate) fn blank_after_summary(checker: &Checker, docstring: &Docstring) {
     let body = docstring.body();
 
-    if !docstring.triple_quoted() {
+    if !docstring.is_triple_quoted() {
         return;
     }
 

crates/ruff_linter/src/rules/pydocstyle/rules/blank_before_after_class.rs

@@ -3,7 +3,7 @@ use ruff_macros::{derive_message_formats, ViolationMetadata};
 use ruff_python_trivia::{indentation_at_offset, PythonWhitespace};
 use ruff_source_file::{Line, LineRanges, UniversalNewlineIterator};
 use ruff_text_size::Ranged;
-use ruff_text_size::{TextLen, TextRange};
+use ruff_text_size::TextRange;
 
 use crate::checkers::ast::Checker;
 use crate::docstrings::Docstring;
@@ -197,7 +197,7 @@ pub(crate) fn blank_before_after_class(checker: &Checker, docstring: &Docstring)
                 // Delete the blank line before the class.
                 diagnostic.set_fix(Fix::safe_edit(Edit::deletion(
                     blank_lines_start,
-                    docstring.start() - docstring.indentation.text_len(),
+                    docstring.line_start(),
                 )));
                 checker.report_diagnostic(diagnostic);
             }
@@ -210,7 +210,7 @@ pub(crate) fn blank_before_after_class(checker: &Checker, docstring: &Docstring)
                 diagnostic.set_fix(Fix::safe_edit(Edit::replacement(
                     checker.stylist().line_ending().to_string(),
                     blank_lines_start,
-                    docstring.start() - docstring.indentation.text_len(),
+                    docstring.line_start(),
                 )));
                 checker.report_diagnostic(diagnostic);
             }

crates/ruff_linter/src/rules/pydocstyle/rules/blank_before_after_function.rs

@@ -6,7 +6,7 @@ use ruff_macros::{derive_message_formats, ViolationMetadata};
 use ruff_python_trivia::PythonWhitespace;
 use ruff_source_file::{UniversalNewlineIterator, UniversalNewlines};
 use ruff_text_size::Ranged;
-use ruff_text_size::{TextLen, TextRange};
+use ruff_text_size::TextRange;
 
 use crate::checkers::ast::Checker;
 use crate::docstrings::Docstring;
@@ -135,7 +135,7 @@ pub(crate) fn blank_before_after_function(checker: &Checker, docstring: &Docstri
             // Delete the blank line before the docstring.
             diagnostic.set_fix(Fix::safe_edit(Edit::deletion(
                 blank_lines_start,
-                docstring.start() - docstring.indentation.text_len(),
+                docstring.line_start(),
             )));
             checker.report_diagnostic(diagnostic);
         }

crates/ruff_linter/src/rules/pydocstyle/rules/indent.rs

@@ -179,8 +179,9 @@ pub(crate) fn indent(checker: &Checker, docstring: &Docstring) {
         return;
     }
 
-    let mut has_seen_tab = docstring.indentation.contains('\t');
-    let docstring_indent_size = docstring.indentation.chars().count();
+    let docstring_indentation = docstring.compute_indentation();
+    let mut has_seen_tab = docstring_indentation.contains('\t');
+    let docstring_indent_size = docstring_indentation.chars().count();
 
     // Lines, other than the last, that are over indented.
     let mut over_indented_lines = vec![];
@@ -226,7 +227,7 @@ pub(crate) fn indent(checker: &Checker, docstring: &Docstring) {
                 let mut diagnostic =
                     Diagnostic::new(UnderIndentation, TextRange::empty(line.start()));
                 diagnostic.set_fix(Fix::safe_edit(Edit::range_replacement(
-                    clean_space(docstring.indentation),
+                    clean_space(docstring_indentation),
                     TextRange::at(line.start(), line_indent.text_len()),
                 )));
                 checker.report_diagnostic(diagnostic);
@@ -275,7 +276,7 @@ pub(crate) fn indent(checker: &Checker, docstring: &Docstring) {
         if let Some(smallest_over_indent_size) = smallest_over_indent_size {
             for line in over_indented_lines {
                 let line_indent = leading_space(&line);
-                let indent = clean_space(docstring.indentation);
+                let indent = clean_space(docstring_indentation);
 
                 // We report over-indentation on every line. This isn't great, but
                 // enables the fix capability.
@@ -324,7 +325,7 @@ pub(crate) fn indent(checker: &Checker, docstring: &Docstring) {
             if last_line_over_indent > 0 && is_indent_only {
                 let mut diagnostic =
                     Diagnostic::new(OverIndentation, TextRange::empty(last.start()));
-                let indent = clean_space(docstring.indentation);
+                let indent = clean_space(docstring_indentation);
                 let range = TextRange::at(last.start(), line_indent.text_len());
                 let edit = if indent.is_empty() {
                     Edit::range_deletion(range)

crates/ruff_linter/src/rules/pydocstyle/rules/multi_line_summary_start.rs

@@ -1,6 +1,8 @@
+use std::borrow::Cow;
+
 use ruff_diagnostics::{AlwaysFixableViolation, Diagnostic, Edit, Fix};
 use ruff_macros::{derive_message_formats, ViolationMetadata};
-use ruff_python_ast::str::{is_triple_quote, leading_quote};
+use ruff_python_ast::str::is_triple_quote;
 use ruff_python_semantic::Definition;
 use ruff_source_file::{LineRanges, NewlineWithTrailingNewline, UniversalNewlineIterator};
 use ruff_text_size::{Ranged, TextRange, TextSize};
@@ -137,7 +139,6 @@ impl AlwaysFixableViolation for MultiLineSummarySecondLine {
 
 /// D212, D213
 pub(crate) fn multi_line_summary_start(checker: &Checker, docstring: &Docstring) {
-    let contents = docstring.contents;
     let body = docstring.body();
 
     if NewlineWithTrailingNewline::from(body.as_str())
@@ -146,7 +147,8 @@ pub(crate) fn multi_line_summary_start(checker: &Checker, docstring: &Docstring)
     {
         return;
     };
-    let mut content_lines = UniversalNewlineIterator::with_offset(contents, docstring.start());
+    let mut content_lines =
+        UniversalNewlineIterator::with_offset(docstring.contents(), docstring.start());
 
     let Some(first_line) = content_lines.next() else {
         return;
@@ -179,7 +181,7 @@ pub(crate) fn multi_line_summary_start(checker: &Checker, docstring: &Docstring)
     } else {
         if checker.enabled(Rule::MultiLineSummarySecondLine) {
             let mut diagnostic = Diagnostic::new(MultiLineSummarySecondLine, docstring.range());
-            let mut indentation = String::from(docstring.indentation);
+            let mut indentation = Cow::Borrowed(docstring.compute_indentation());
             let mut fixable = true;
             if !indentation.chars().all(char::is_whitespace) {
                 fixable = false;
@@ -193,6 +195,7 @@ pub(crate) fn multi_line_summary_start(checker: &Checker, docstring: &Docstring)
                         .slice(TextRange::new(stmt_line_start, member.start()));
 
                     if stmt_indentation.chars().all(char::is_whitespace) {
+                        let indentation = indentation.to_mut();
                         indentation.clear();
                         indentation.push_str(stmt_indentation);
                         indentation.push_str(checker.stylist().indentation());
@@ -202,14 +205,16 @@ pub(crate) fn multi_line_summary_start(checker: &Checker, docstring: &Docstring)
             }
 
             if fixable {
-                let prefix = leading_quote(contents).unwrap();
                 // Use replacement instead of insert to trim possible whitespace between leading
                 // quote and text.
                 let repl = format!(
                     "{}{}{}",
                     checker.stylist().line_ending().as_str(),
                     indentation,
-                    first_line.strip_prefix(prefix).unwrap().trim_start()
+                    first_line
+                        .strip_prefix(docstring.opener())
+                        .unwrap()
+                        .trim_start()
                 );
 
                 diagnostic.set_fix(Fix::safe_edit(Edit::replacement(

crates/ruff_linter/src/rules/pydocstyle/rules/newline_after_last_paragraph.rs

@@ -59,10 +59,10 @@ impl AlwaysFixableViolation for NewLineAfterLastParagraph {
 
 /// D209
 pub(crate) fn newline_after_last_paragraph(checker: &Checker, docstring: &Docstring) {
-    let contents = docstring.contents;
+    let contents = docstring.contents();
     let body = docstring.body();
 
-    if !docstring.triple_quoted() {
+    if !docstring.is_triple_quoted() {
         return;
     }
 
@@ -92,7 +92,7 @@ pub(crate) fn newline_after_last_paragraph(checker: &Checker, docstring: &Docstr
                     let content = format!(
                         "{}{}",
                         checker.stylist().line_ending().as_str(),
-                        clean_space(docstring.indentation)
+                        clean_space(docstring.compute_indentation())
                     );
                     diagnostic.set_fix(Fix::safe_edit(Edit::replacement(
                         content,

crates/ruff_linter/src/rules/pydocstyle/rules/no_surrounding_whitespace.rs

@@ -63,7 +63,7 @@ pub(crate) fn no_surrounding_whitespace(checker: &Checker, docstring: &Docstring
         return;
     }
     let mut diagnostic = Diagnostic::new(SurroundingWhitespace, docstring.range());
-    let quote = docstring.contents.chars().last().unwrap();
+    let quote = docstring.quote_style().as_char();
     // If removing whitespace would lead to an invalid string of quote
     // characters, avoid applying the fix.
     if !trimmed.ends_with(quote) && !trimmed.starts_with(quote) && !ends_with_backslash(trimmed) {

crates/ruff_linter/src/rules/pydocstyle/rules/one_liner.rs

@@ -1,6 +1,5 @@
 use ruff_diagnostics::{Diagnostic, Edit, Fix, FixAvailability, Violation};
 use ruff_macros::{derive_message_formats, ViolationMetadata};
-use ruff_python_ast::str::{leading_quote, trailing_quote};
 use ruff_source_file::NewlineWithTrailingNewline;
 use ruff_text_size::Ranged;
 
@@ -64,24 +63,26 @@ pub(crate) fn one_liner(checker: &Checker, docstring: &Docstring) {
 
     if non_empty_line_count == 1 && line_count > 1 {
         let mut diagnostic = Diagnostic::new(UnnecessaryMultilineDocstring, docstring.range());
-        if let (Some(leading), Some(trailing)) = (
-            leading_quote(docstring.contents),
-            trailing_quote(docstring.contents),
-        ) {
+
         // If removing whitespace would lead to an invalid string of quote
         // characters, avoid applying the fix.
         let body = docstring.body();
         let trimmed = body.trim();
+        let quote_char = docstring.quote_style().as_char();
         if trimmed.chars().rev().take_while(|c| *c == '\\').count() % 2 == 0
-                && !trimmed.ends_with(trailing.chars().last().unwrap())
-                && !trimmed.starts_with(leading.chars().last().unwrap())
+            && !trimmed.ends_with(quote_char)
+            && !trimmed.starts_with(quote_char)
         {
             diagnostic.set_fix(Fix::unsafe_edit(Edit::range_replacement(
-                    format!("{leading}{trimmed}{trailing}"),
+                format!(
+                    "{leading}{trimmed}{trailing}",
+                    leading = docstring.opener(),
+                    trailing = docstring.closer()
+                ),
                 docstring.range(),
             )));
         }
-        }
+
         checker.report_diagnostic(diagnostic);
     }
 }

crates/ruff_linter/src/rules/pydocstyle/rules/sections.rs

@@ -1399,7 +1399,8 @@ fn blanks_and_section_underline(
 
             if checker.enabled(Rule::OverindentedSectionUnderline) {
                 let leading_space = leading_space(&non_blank_line);
-                if leading_space.len() > docstring.indentation.len() {
+                let docstring_indentation = docstring.compute_indentation();
+                if leading_space.len() > docstring_indentation.len() {
                     let mut diagnostic = Diagnostic::new(
                         OverindentedSectionUnderline {
                             name: context.section_name().to_string(),
@@ -1412,7 +1413,7 @@ fn blanks_and_section_underline(
                         blank_lines_end,
                         leading_space.text_len() + TextSize::from(1),
                     );
-                    let contents = clean_space(docstring.indentation);
+                    let contents = clean_space(docstring_indentation);
                     diagnostic.set_fix(Fix::safe_edit(if contents.is_empty() {
                         Edit::range_deletion(range)
                     } else {
@@ -1540,7 +1541,7 @@ fn blanks_and_section_underline(
                     let content = format!(
                         "{}{}{}",
                         checker.stylist().line_ending().as_str(),
-                        clean_space(docstring.indentation),
+                        clean_space(docstring.compute_indentation()),
                         "-".repeat(context.section_name().len()),
                     );
                     diagnostic.set_fix(Fix::safe_edit(Edit::insertion(
@@ -1621,7 +1622,7 @@ fn blanks_and_section_underline(
             let content = format!(
                 "{}{}{}",
                 checker.stylist().line_ending().as_str(),
-                clean_space(docstring.indentation),
+                clean_space(docstring.compute_indentation()),
                 "-".repeat(context.section_name().len()),
             );
             diagnostic.set_fix(Fix::safe_edit(Edit::insertion(
@@ -1671,7 +1672,8 @@ fn common_section(
 
     if checker.enabled(Rule::OverindentedSection) {
         let leading_space = leading_space(context.summary_line());
-        if leading_space.len() > docstring.indentation.len() {
+        let docstring_indentation = docstring.compute_indentation();
+        if leading_space.len() > docstring_indentation.len() {
             let section_range = context.section_name_range();
             let mut diagnostic = Diagnostic::new(
                 OverindentedSection {
@@ -1681,7 +1683,7 @@ fn common_section(
             );
 
             // Replace the existing indentation with whitespace of the appropriate length.
-            let content = clean_space(docstring.indentation);
+            let content = clean_space(docstring_indentation);
             let fix_range = TextRange::at(context.start(), leading_space.text_len());
             diagnostic.set_fix(Fix::safe_edit(if content.is_empty() {
                 Edit::range_deletion(fix_range)
@@ -1738,7 +1740,7 @@ fn common_section(
                     format!(
                         "{}{}",
                         line_end.repeat(2 - num_blank_lines),
-                        docstring.indentation
+                        docstring.compute_indentation()
                     ),
                     context.end() - del_len,
                     context.end(),

crates/ruff_linter/src/rules/pydocstyle/rules/triple_quotes.rs

@@ -64,9 +64,8 @@ impl Violation for TripleSingleQuotes {
 
 /// D300
 pub(crate) fn triple_quotes(checker: &Checker, docstring: &Docstring) {
-    let leading_quote = docstring.leading_quote();
-
-    let prefixes = leading_quote.trim_end_matches(['\'', '"']).to_owned();
+    let opener = docstring.opener();
+    let prefixes = docstring.prefix_str();
 
     let expected_quote = if docstring.body().contains("\"\"\"") {
         if docstring.body().contains("\'\'\'") {
@@ -79,7 +78,7 @@ pub(crate) fn triple_quotes(checker: &Checker, docstring: &Docstring) {
 
     match expected_quote {
         Quote::Single => {
-            if !leading_quote.ends_with("'''") {
+            if !opener.ends_with("'''") {
                 let mut diagnostic =
                     Diagnostic::new(TripleSingleQuotes { expected_quote }, docstring.range());
 
@@ -95,7 +94,7 @@ pub(crate) fn triple_quotes(checker: &Checker, docstring: &Docstring) {
             }
         }
         Quote::Double => {
-            if !leading_quote.ends_with("\"\"\"") {
+            if !opener.ends_with("\"\"\"") {
                 let mut diagnostic =
                     Diagnostic::new(TripleSingleQuotes { expected_quote }, docstring.range());
 

crates/ruff_python_ast/src/nodes.rs

@@ -1645,6 +1645,16 @@ impl StringLiteral {
             flags: StringLiteralFlags::empty().with_invalid(),
         }
     }
+
+    /// The range of the string literal's contents.
+    ///
+    /// This excludes any prefixes, opening quotes or closing quotes.
+    pub fn content_range(&self) -> TextRange {
+        TextRange::new(
+            self.start() + self.flags.opener_len(),
+            self.end() - self.flags.closer_len(),
+        )
+    }
 }
 
 impl From<StringLiteral> for Expr {

crates/ruff_python_ast/src/str_prefix.rs

@@ -1,3 +1,5 @@
+use ruff_text_size::TextSize;
+
 use std::fmt;
 
 /// Enumerations of the valid prefixes a string literal can have.
@@ -33,6 +35,13 @@ impl StringLiteralPrefix {
             Self::Raw { uppercase: false } => "r",
         }
     }
+
+    pub const fn text_len(self) -> TextSize {
+        match self {
+            Self::Empty => TextSize::new(0),
+            Self::Unicode | Self::Raw { .. } => TextSize::new(1),
+        }
+    }
 }
 
 impl fmt::Display for StringLiteralPrefix {