mirrors/ruff
1
0
Fork 0
mirror of https://github.com/astral-sh/ruff.git synced 2025-07-23 04:55:09 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 4

fmt: off..on suppression comments (#6477)

Browse source
This commit is contained in:
Micha Reiser 2023-08-14 17:57:36 +02:00 • committed by GitHub
parent 278a4f6e14
commit 09c8b17661
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
34 changed files with 1883 additions and 978 deletions
Download patch file Download diff file Expand all files Collapse all files

8
crates/ruff_python_formatter/src/comments/format.rs
View file

@ -281,11 +281,11 @@ impl Format<PyFormatContext<'_>> for FormatDanglingOpenParenthesisComments<'_> {
///
/// * Adds a whitespace between `#` and the comment text except if the first character is a `#`, `:`, `'`, or `!`
/// * Replaces non breaking whitespaces with regular whitespaces except if in front of a `types:` comment
const fn format_comment(comment: &SourceComment) -> FormatComment {
pub(crate) const fn format_comment(comment: &SourceComment) -> FormatComment {
FormatComment { comment }
}
struct FormatComment<'a> {
pub(crate) struct FormatComment<'a> {
comment: &'a SourceComment,
}
@ -343,12 +343,12 @@ impl Format<PyFormatContext<'_>> for FormatComment<'_> {
// Top level: Up to two empty lines
// parenthesized: A single empty line
// other: Up to a single empty line
const fn empty_lines(lines: u32) -> FormatEmptyLines {
pub(crate) const fn empty_lines(lines: u32) -> FormatEmptyLines {
FormatEmptyLines { lines }
}
#[derive(Copy, Clone, Debug)]
struct FormatEmptyLines {
pub(crate) struct FormatEmptyLines {
lines: u32,
}

1
crates/ruff_python_formatter/src/comments/map.rs
View file

@ -244,6 +244,7 @@ impl<K: std::hash::Hash + Eq, V> MultiMap<K, V> {
}
/// Returns `true` if `key` has any *leading*, *dangling*, or *trailing* parts.
#[allow(unused)]
pub(super) fn has(&self, key: &K) -> bool {
self.index.get(key).is_some()
}

64
crates/ruff_python_formatter/src/comments/mod.rs
View file

@ -103,6 +103,7 @@ use ruff_formatter::{SourceCode, SourceCodeSlice};
use ruff_python_ast::node::AnyNodeRef;
use ruff_python_ast::visitor::preorder::{PreorderVisitor, TraversalSignal};
use ruff_python_index::CommentRanges;
use ruff_python_trivia::PythonWhitespace;
use crate::comments::debug::{DebugComment, DebugComments};
use crate::comments::map::MultiMap;
@ -110,7 +111,7 @@ use crate::comments::node_key::NodeRefEqualityKey;
use crate::comments::visitor::CommentsVisitor;
mod debug;
mod format;
pub(crate) mod format;
mod map;
mod node_key;
mod placement;
@ -150,6 +151,11 @@ impl SourceComment {
self.formatted.set(true);
}
/// Marks the comment as not-formatted
pub(crate) fn mark_unformatted(&self) {
self.formatted.set(false);
}
/// If the comment has already been formatted
pub(crate) fn is_formatted(&self) -> bool {
self.formatted.get()
@ -163,6 +169,50 @@ impl SourceComment {
pub(crate) fn debug<'a>(&'a self, source_code: SourceCode<'a>) -> DebugComment<'a> {
DebugComment::new(self, source_code)
}
pub(crate) fn suppression_kind(&self, source: &str) -> Option<SuppressionKind> {
let text = self.slice.text(SourceCode::new(source));
let trimmed = text.strip_prefix('#').unwrap_or(text).trim_whitespace();
if let Some(command) = trimmed.strip_prefix("fmt:") {
match command.trim_whitespace_start() {
"off" => Some(SuppressionKind::Off),
"on" => Some(SuppressionKind::On),
"skip" => Some(SuppressionKind::Skip),
_ => None,
}
} else if let Some(command) = trimmed.strip_prefix("yapf:") {
match command.trim_whitespace_start() {
"disable" => Some(SuppressionKind::Off),
"enable" => Some(SuppressionKind::On),
_ => None,
}
} else {
None
}
}
/// Returns true if this comment is a `fmt: off` or `yapf: disable` own line suppression comment.
pub(crate) fn is_suppression_off_comment(&self, source: &str) -> bool {
self.line_position.is_own_line()
&& matches!(self.suppression_kind(source), Some(SuppressionKind::Off))
}
/// Returns true if this comment is a `fmt: on` or `yapf: enable` own line suppression comment.
pub(crate) fn is_suppression_on_comment(&self, source: &str) -> bool {
self.line_position.is_own_line()
&& matches!(self.suppression_kind(source), Some(SuppressionKind::On))
}
}
#[derive(Copy, Clone, Eq, PartialEq, Debug)]
pub(crate) enum SuppressionKind {
/// A `fmt: off` or `yapf: disable` comment
Off,
/// A `fmt: on` or `yapf: enable` comment
On,
/// A `fmt: skip` comment
Skip,
}
impl Ranged for SourceComment {
@ -246,8 +296,6 @@ pub(crate) struct Comments<'a> {
data: Rc<CommentsData<'a>>,
}
#[allow(unused)]
// TODO(micha): Remove after using the new comments infrastructure in the formatter.
impl<'a> Comments<'a> {
fn new(comments: CommentsMap<'a>) -> Self {
Self {
@ -270,16 +318,6 @@ impl<'a> Comments<'a> {
Self::new(map)
}
#[inline]
pub(crate) fn has_comments<T>(&self, node: T) -> bool
where
T: Into<AnyNodeRef<'a>>,
{
self.data
.comments
.has(&NodeRefEqualityKey::from_ref(node.into()))
}
/// Returns `true` if the given `node` has any [leading comments](self#leading-comments).
#[inline]
pub(crate) fn has_leading_comments<T>(&self, node: T) -> bool

7
crates/ruff_python_formatter/src/expression/expr_ipy_escape_command.rs
View file

@ -1,12 +1,11 @@
use crate::{verbatim_text, FormatNodeRule, PyFormatter};
use ruff_formatter::{write, Buffer, FormatResult};
use ruff_python_ast::ExprIpyEscapeCommand;
use crate::prelude::*;
use ruff_python_ast::{ExprIpyEscapeCommand, Ranged};
#[derive(Default)]
pub struct FormatExprIpyEscapeCommand;
impl FormatNodeRule<ExprIpyEscapeCommand> for FormatExprIpyEscapeCommand {
fn fmt_fields(&self, item: &ExprIpyEscapeCommand, f: &mut PyFormatter) -> FormatResult<()> {
write!(f, [verbatim_text(item)])
source_text_slice(item.range(), ContainsNewlines::No).fmt(f)
}
}

100
crates/ruff_python_formatter/src/lib.rs
View file

@ -1,26 +1,24 @@
use thiserror::Error;
use ruff_formatter::format_element::tag;
use ruff_formatter::prelude::{source_position, text, Formatter, Tag};
use ruff_formatter::{
format, write, Buffer, Format, FormatElement, FormatError, FormatResult, PrintError,
};
use ruff_formatter::{Formatted, Printed, SourceCode};
use ruff_python_ast::node::{AnyNodeRef, AstNode};
use ruff_python_ast::Mod;
use ruff_python_index::{CommentRanges, CommentRangesBuilder};
use ruff_python_parser::lexer::{lex, LexicalError};
use ruff_python_parser::{parse_tokens, Mode, ParseError};
use ruff_source_file::Locator;
use ruff_text_size::TextLen;
use crate::comments::{
dangling_node_comments, leading_node_comments, trailing_node_comments, Comments,
};
use crate::context::PyFormatContext;
pub use crate::options::{MagicTrailingComma, PyFormatOptions, QuoteStyle};
use ruff_formatter::format_element::tag;
use ruff_formatter::prelude::{
dynamic_text, source_position, source_text_slice, text, ContainsNewlines, Formatter, Tag,
};
use ruff_formatter::{
format, normalize_newlines, write, Buffer, Format, FormatElement, FormatError, FormatResult,
PrintError,
};
use ruff_formatter::{Formatted, Printed, SourceCode};
use ruff_python_ast::node::{AnyNodeRef, AstNode};
use ruff_python_ast::{Mod, Ranged};
use ruff_python_index::{CommentRanges, CommentRangesBuilder};
use ruff_python_parser::lexer::{lex, LexicalError};
use ruff_python_parser::{parse_tokens, Mode, ParseError};
use ruff_source_file::Locator;
use ruff_text_size::{TextLen, TextRange};
use std::borrow::Cow;
use thiserror::Error;
pub(crate) mod builders;
pub mod cli;
@ -35,6 +33,7 @@ pub(crate) mod pattern;
mod prelude;
pub(crate) mod statement;
pub(crate) mod type_param;
mod verbatim;
include!("../../ruff_formatter/shared_traits.rs");
@ -47,10 +46,10 @@ where
N: AstNode,
{
fn fmt(&self, node: &N, f: &mut PyFormatter) -> FormatResult<()> {
self.fmt_leading_comments(node, f)?;
leading_node_comments(node).fmt(f)?;
self.fmt_node(node, f)?;
self.fmt_dangling_comments(node, f)?;
self.fmt_trailing_comments(node, f)
trailing_node_comments(node).fmt(f)
}
/// Formats the node without comments. Ignores any suppression comments.
@ -63,14 +62,6 @@ where
/// Formats the node's fields.
fn fmt_fields(&self, item: &N, f: &mut PyFormatter) -> FormatResult<()>;
/// Formats the [leading comments](comments#leading-comments) of the node.
///
/// You may want to override this method if you want to manually handle the formatting of comments
/// inside of the `fmt_fields` method or customize the formatting of the leading comments.
fn fmt_leading_comments(&self, node: &N, f: &mut PyFormatter) -> FormatResult<()> {
leading_node_comments(node).fmt(f)
}
/// Formats the [dangling comments](comments#dangling-comments) of the node.
///
/// You should override this method if the node handled by this rule can have dangling comments because the
@ -81,14 +72,6 @@ where
fn fmt_dangling_comments(&self, node: &N, f: &mut PyFormatter) -> FormatResult<()> {
dangling_node_comments(node).fmt(f)
}
/// Formats the [trailing comments](comments#trailing-comments) of the node.
///
/// You may want to override this method if you want to manually handle the formatting of comments
/// inside of the `fmt_fields` method or customize the formatting of the trailing comments.
fn fmt_trailing_comments(&self, node: &N, f: &mut PyFormatter) -> FormatResult<()> {
trailing_node_comments(node).fmt(f)
}
}
#[derive(Error, Debug)]
@ -234,53 +217,18 @@ impl Format<PyFormatContext<'_>> for NotYetImplementedCustomText<'_> {
}
}
pub(crate) struct VerbatimText(TextRange);
#[allow(unused)]
pub(crate) fn verbatim_text<T>(item: &T) -> VerbatimText
where
T: Ranged,
{
VerbatimText(item.range())
}
impl Format<PyFormatContext<'_>> for VerbatimText {
fn fmt(&self, f: &mut PyFormatter) -> FormatResult<()> {
f.write_element(FormatElement::Tag(Tag::StartVerbatim(
tag::VerbatimKind::Verbatim {
length: self.0.len(),
},
)))?;
match normalize_newlines(f.context().locator().slice(self.0), ['\r']) {
Cow::Borrowed(_) => {
write!(f, [source_text_slice(self.0, ContainsNewlines::Detect)])?;
}
Cow::Owned(cleaned) => {
write!(
f,
[
dynamic_text(&cleaned, Some(self.0.start())),
source_position(self.0.end())
]
)?;
}
}
f.write_element(FormatElement::Tag(Tag::EndVerbatim))?;
Ok(())
}
}
#[cfg(test)]
mod tests {
use crate::{format_module, format_node, PyFormatOptions};
use std::path::Path;
use anyhow::Result;
use insta::assert_snapshot;
use ruff_python_index::CommentRangesBuilder;
use ruff_python_parser::lexer::lex;
use ruff_python_parser::{parse_tokens, Mode};
use std::path::Path;
use crate::{format_module, format_node, PyFormatOptions};
/// Very basic test intentionally kept very similar to the CLI
#[test]

7
crates/ruff_python_formatter/src/statement/stmt_ipy_escape_command.rs
View file

@ -1,12 +1,11 @@
use crate::{verbatim_text, FormatNodeRule, PyFormatter};
use ruff_formatter::{write, Buffer, FormatResult};
use ruff_python_ast::StmtIpyEscapeCommand;
use crate::prelude::*;
use ruff_python_ast::{Ranged, StmtIpyEscapeCommand};
#[derive(Default)]
pub struct FormatStmtIpyEscapeCommand;
impl FormatNodeRule<StmtIpyEscapeCommand> for FormatStmtIpyEscapeCommand {
fn fmt_fields(&self, item: &StmtIpyEscapeCommand, f: &mut PyFormatter) -> FormatResult<()> {
write!(f, [verbatim_text(item)])
source_text_slice(item.range(), ContainsNewlines::No).fmt(f)
}
}

498
crates/ruff_python_formatter/src/statement/suite.rs
View file

@ -1,14 +1,19 @@
use crate::comments::{leading_comments, trailing_comments};
use ruff_formatter::{write, FormatOwnedWithRule, FormatRefWithRule, FormatRuleWithOptions};
use ruff_python_ast::helpers::is_compound_statement;
use ruff_python_ast::{self as ast, Ranged, Stmt, Suite};
use ruff_python_ast::{Constant, ExprConstant};
use ruff_python_ast::node::AnyNodeRef;
use ruff_python_ast::{self as ast, Expr, ExprConstant, Ranged, Stmt, Suite};
use ruff_python_trivia::{lines_after_ignoring_trivia, lines_before};
use ruff_text_size::TextRange;
use crate::comments::{leading_comments, trailing_comments};
use crate::context::{NodeLevel, WithNodeLevel};
use crate::expression::expr_constant::ExprConstantLayout;
use crate::expression::string::StringLayout;
use crate::prelude::*;
use crate::verbatim::{
write_suppressed_statements_starting_with_leading_comment,
write_suppressed_statements_starting_with_trailing_comment,
};
/// Level at which the [`Suite`] appears in the source code.
#[derive(Copy, Clone, Debug)]
@ -51,196 +56,231 @@ impl FormatRule<Suite, PyFormatContext<'_>> for FormatSuite {
let comments = f.context().comments().clone();
let source = f.context().source();
let mut iter = statements.iter();
let Some(first) = iter.next() else {
return Ok(());
};
let mut f = WithNodeLevel::new(node_level, f);
write!(
f,
[format_with(|f| {
let mut iter = statements.iter();
let Some(first) = iter.next() else {
return Ok(());
};
// Format the first statement in the body, which often has special formatting rules.
let mut last = first;
match self.kind {
SuiteKind::Other => {
if is_class_or_function_definition(first) && !comments.has_leading_comments(first) {
// Add an empty line for any nested functions or classes defined within
// non-function or class compound statements, e.g., this is stable formatting:
// ```python
// if True:
//
// def test():
// ...
// ```
write!(f, [empty_line()])?;
}
write!(f, [first.format()])?;
}
SuiteKind::Function => {
if let Some(constant) = get_docstring(first) {
write!(
f,
[
// We format the expression, but the statement carries the comments
leading_comments(comments.leading_comments(first)),
constant
.format()
.with_options(ExprConstantLayout::String(StringLayout::DocString)),
trailing_comments(comments.trailing_comments(first)),
]
)?;
// Format the first statement in the body, which often has special formatting rules.
let first = match self.kind {
SuiteKind::Other => {
if is_class_or_function_definition(first)
&& !comments.has_leading_comments(first)
{
// Add an empty line for any nested functions or classes defined within
// non-function or class compound statements, e.g., this is stable formatting:
// ```python
// if True:
//
// def test():
// ...
// ```
empty_line().fmt(f)?;
}
SuiteChildStatement::Other(first)
}
SuiteKind::Function => {
if let Some(docstring) = DocstringStmt::try_from_statement(first) {
SuiteChildStatement::Docstring(docstring)
} else {
SuiteChildStatement::Other(first)
}
}
SuiteKind::Class => {
if let Some(docstring) = DocstringStmt::try_from_statement(first) {
if !comments.has_leading_comments(first)
&& lines_before(first.start(), source) > 1
{
// Allow up to one empty line before a class docstring, e.g., this is
// stable formatting:
// ```python
// class Test:
//
// """Docstring"""
// ```
empty_line().fmt(f)?;
}
SuiteChildStatement::Docstring(docstring)
} else {
SuiteChildStatement::Other(first)
}
}
SuiteKind::TopLevel => SuiteChildStatement::Other(first),
};
let (mut preceding, mut after_class_docstring) = if comments
.leading_comments(first)
.iter()
.any(|comment| comment.is_suppression_off_comment(source))
{
(
write_suppressed_statements_starting_with_leading_comment(
first, &mut iter, f,
)?,
false,
)
} else if comments
.trailing_comments(first)
.iter()
.any(|comment| comment.is_suppression_off_comment(source))
{
(
write_suppressed_statements_starting_with_trailing_comment(
first, &mut iter, f,
)?,
false,
)
} else {
write!(f, [first.format()])?;
}
}
SuiteKind::Class => {
if let Some(constant) = get_docstring(first) {
if !comments.has_leading_comments(first)
&& lines_before(first.start(), source) > 1
first.fmt(f)?;
(
first.statement(),
matches!(first, SuiteChildStatement::Docstring(_))
&& matches!(self.kind, SuiteKind::Class),
)
};
while let Some(following) = iter.next() {
if is_class_or_function_definition(preceding)
|| is_class_or_function_definition(following)
{
// Allow up to one empty line before a class docstring
match self.kind {
SuiteKind::TopLevel => {
write!(f, [empty_line(), empty_line()])?;
}
SuiteKind::Function | SuiteKind::Class | SuiteKind::Other => {
empty_line().fmt(f)?;
}
}
} else if is_import_definition(preceding) && !is_import_definition(following) {
empty_line().fmt(f)?;
} else if is_compound_statement(preceding) {
// Handles the case where a body has trailing comments. The issue is that RustPython does not include
// the comments in the range of the suite. This means, the body ends right after the last statement in the body.
// ```python
// def test():
// ...
// # The body of `test` ends right after `...` and before this comment
//
// # leading comment
//
//
// a = 10
// ```
// Using `lines_after` for the node doesn't work because it would count the lines after the `...`
// which is 0 instead of 1, the number of lines between the trailing comment and
// the leading comment. This is why the suite handling counts the lines before the
// start of the next statement or before the first leading comments for compound statements.
let start = if let Some(first_leading) =
comments.leading_comments(following).first()
{
first_leading.slice().start()
} else {
following.start()
};
match lines_before(start, source) {
0 | 1 => hard_line_break().fmt(f)?,
2 => empty_line().fmt(f)?,
3.. => match self.kind {
SuiteKind::TopLevel => write!(f, [empty_line(), empty_line()])?,
SuiteKind::Function | SuiteKind::Class | SuiteKind::Other => {
empty_line().fmt(f)?;
}
},
}
} else if after_class_docstring {
// Enforce an empty line after a class docstring, e.g., these are both stable
// formatting:
// ```python
// class Test:
// """Docstring"""
//
// ...
//
//
// class Test:
//
// """Docstring"""
//
// ...
// ```
write!(f, [empty_line()])?;
}
write!(
f,
[
// We format the expression, but the statement carries the comments
leading_comments(comments.leading_comments(first)),
constant
.format()
.with_options(ExprConstantLayout::String(StringLayout::DocString)),
trailing_comments(comments.trailing_comments(first)),
]
)?;
// Enforce an empty line after a class docstring
// ```python
// class Test:
// """Docstring"""
//
// ...
//
//
// class Test:
//
// """Docstring"""
//
// ...
// ```
// Unlike black, we add the newline also after single quoted docstrings
if let Some(second) = iter.next() {
// Format the subsequent statement immediately. This rule takes precedence
// over the rules in the loop below (and most of them won't apply anyway,
// e.g., we know the first statement isn't an import).
write!(f, [empty_line(), second.format()])?;
last = second;
}
} else {
// No docstring, use normal formatting
write!(f, [first.format()])?;
}
}
SuiteKind::TopLevel => {
write!(f, [first.format()])?;
}
}
for statement in iter {
if is_class_or_function_definition(last) || is_class_or_function_definition(statement) {
match self.kind {
SuiteKind::TopLevel => {
write!(f, [empty_line(), empty_line(), statement.format()])?;
}
SuiteKind::Function | SuiteKind::Class | SuiteKind::Other => {
write!(f, [empty_line(), statement.format()])?;
}
}
} else if is_import_definition(last) && !is_import_definition(statement) {
write!(f, [empty_line(), statement.format()])?;
} else if is_compound_statement(last) {
// Handles the case where a body has trailing comments. The issue is that RustPython does not include
// the comments in the range of the suite. This means, the body ends right after the last statement in the body.
// ```python
// def test():
// ...
// # The body of `test` ends right after `...` and before this comment
//
// # leading comment
//
//
// a = 10
// ```
// Using `lines_after` for the node doesn't work because it would count the lines after the `...`
// which is 0 instead of 1, the number of lines between the trailing comment and
// the leading comment. This is why the suite handling counts the lines before the
// start of the next statement or before the first leading comments for compound statements.
let start =
if let Some(first_leading) = comments.leading_comments(statement).first() {
first_leading.slice().start()
empty_line().fmt(f)?;
after_class_docstring = false;
} else {
statement.start()
};
// Insert the appropriate number of empty lines based on the node level, e.g.:
// * [`NodeLevel::Module`]: Up to two empty lines
// * [`NodeLevel::CompoundStatement`]: Up to one empty line
// * [`NodeLevel::Expression`]: No empty lines
match lines_before(start, source) {
0 | 1 => write!(f, [hard_line_break()])?,
2 => write!(f, [empty_line()])?,
3.. => match self.kind {
SuiteKind::TopLevel => write!(f, [empty_line(), empty_line()])?,
SuiteKind::Function | SuiteKind::Class | SuiteKind::Other => {
write!(f, [empty_line()])?;
let count_lines = |offset| {
// It's necessary to skip any trailing line comment because RustPython doesn't include trailing comments
// in the node's range
// ```python
// a # The range of `a` ends right before this comment
//
// b
// ```
//
// Simply using `lines_after` doesn't work if a statement has a trailing comment because
// it then counts the lines between the statement and the trailing comment, which is
// always 0. This is why it skips any trailing trivia (trivia that's on the same line)
// and counts the lines after.
lines_after_ignoring_trivia(offset, source)
};
match node_level {
NodeLevel::TopLevel => match count_lines(preceding.end()) {
0 | 1 => hard_line_break().fmt(f)?,
2 => empty_line().fmt(f)?,
_ => write!(f, [empty_line(), empty_line()])?,
},
NodeLevel::CompoundStatement => match count_lines(preceding.end()) {
0 | 1 => hard_line_break().fmt(f)?,
_ => empty_line().fmt(f)?,
},
NodeLevel::Expression(_) | NodeLevel::ParenthesizedExpression => {
hard_line_break().fmt(f)?;
}
}
},
}
}
write!(f, [statement.format()])?;
} else {
// Insert the appropriate number of empty lines based on the node level, e.g.:
// * [`NodeLevel::Module`]: Up to two empty lines
// * [`NodeLevel::CompoundStatement`]: Up to one empty line
// * [`NodeLevel::Expression`]: No empty lines
let count_lines = |offset| {
// It's necessary to skip any trailing line comment because RustPython doesn't include trailing comments
// in the node's range
// ```python
// a # The range of `a` ends right before this comment
//
// b
// ```
//
// Simply using `lines_after` doesn't work if a statement has a trailing comment because
// it then counts the lines between the statement and the trailing comment, which is
// always 0. This is why it skips any trailing trivia (trivia that's on the same line)
// and counts the lines after.
lines_after_ignoring_trivia(offset, source)
};
match node_level {
NodeLevel::TopLevel => match count_lines(last.end()) {
0 | 1 => write!(f, [hard_line_break()])?,
2 => write!(f, [empty_line()])?,
_ => write!(f, [empty_line(), empty_line()])?,
},
NodeLevel::CompoundStatement => match count_lines(last.end()) {
0 | 1 => write!(f, [hard_line_break()])?,
_ => write!(f, [empty_line()])?,
},
NodeLevel::Expression(_) | NodeLevel::ParenthesizedExpression => {
write!(f, [hard_line_break()])?;
if comments
.leading_comments(following)
.iter()
.any(|comment| comment.is_suppression_off_comment(source))
{
preceding = write_suppressed_statements_starting_with_leading_comment(
SuiteChildStatement::Other(following),
&mut iter,
f,
)?;
} else if comments
.trailing_comments(following)
.iter()
.any(|comment| comment.is_suppression_off_comment(source))
{
preceding = write_suppressed_statements_starting_with_trailing_comment(
SuiteChildStatement::Other(following),
&mut iter,
f,
)?;
} else {
following.format().fmt(f)?;
preceding = following;
}
}
write!(f, [statement.format()])?;
}
last = statement;
}
Ok(())
Ok(())
})]
)
}
}
@ -254,23 +294,6 @@ const fn is_import_definition(stmt: &Stmt) -> bool {
matches!(stmt, Stmt::Import(_) | Stmt::ImportFrom(_))
}
/// Checks if the statement is a simple string that can be formatted as a docstring
fn get_docstring(stmt: &Stmt) -> Option<&ExprConstant> {
let stmt_expr = stmt.as_expr_stmt()?;
let expr_constant = stmt_expr.value.as_constant_expr()?;
if matches!(
expr_constant.value,
Constant::Str(ast::StringConstant {
implicit_concatenated: false,
..
})
) {
Some(expr_constant)
} else {
None
}
}
impl FormatRuleWithOptions<Suite, PyFormatContext<'_>> for FormatSuite {
type Options = SuiteKind;
@ -296,6 +319,93 @@ impl<'ast> IntoFormat<PyFormatContext<'ast>> for Suite {
}
}
/// A statement representing a docstring.
#[derive(Copy, Clone)]
pub(crate) struct DocstringStmt<'a>(&'a Stmt);
impl<'a> DocstringStmt<'a> {
/// Checks if the statement is a simple string that can be formatted as a docstring
fn try_from_statement(stmt: &'a Stmt) -> Option<DocstringStmt<'a>> {
let Stmt::Expr(ast::StmtExpr { value, .. }) = stmt else {
return None;
};
if let Expr::Constant(ExprConstant { value, .. }) = value.as_ref() {
if !value.is_implicit_concatenated() {
return Some(DocstringStmt(stmt));
}
}
None
}
}
impl Format<PyFormatContext<'_>> for DocstringStmt<'_> {
fn fmt(&self, f: &mut Formatter<PyFormatContext<'_>>) -> FormatResult<()> {
// SAFETY: Safe because `DocStringStmt` guarantees that it only ever wraps a `ExprStmt` containing a `ConstantExpr`.
let constant = self
.0
.as_expr_stmt()
.unwrap()
.value
.as_constant_expr()
.unwrap();
let comments = f.context().comments().clone();
// We format the expression, but the statement carries the comments
write!(
f,
[
leading_comments(comments.leading_comments(self.0)),
constant
.format()
.with_options(ExprConstantLayout::String(StringLayout::DocString)),
trailing_comments(comments.trailing_comments(self.0)),
]
)
}
}
/// A Child of a suite.
#[derive(Copy, Clone)]
pub(crate) enum SuiteChildStatement<'a> {
/// A docstring documenting a class or function definition.
Docstring(DocstringStmt<'a>),
/// Any other statement.
Other(&'a Stmt),
}
impl<'a> SuiteChildStatement<'a> {
pub(crate) const fn statement(self) -> &'a Stmt {
match self {
SuiteChildStatement::Docstring(docstring) => docstring.0,
SuiteChildStatement::Other(statement) => statement,
}
}
}
impl Ranged for SuiteChildStatement<'_> {
fn range(&self) -> TextRange {
self.statement().range()
}
}
impl<'a> From<SuiteChildStatement<'a>> for AnyNodeRef<'a> {
fn from(value: SuiteChildStatement<'a>) -> Self {
value.statement().into()
}
}
impl Format<PyFormatContext<'_>> for SuiteChildStatement<'_> {
fn fmt(&self, f: &mut Formatter<PyFormatContext<'_>>) -> FormatResult<()> {
match self {
SuiteChildStatement::Docstring(docstring) => docstring.fmt(f),
SuiteChildStatement::Other(statement) => statement.format().fmt(f),
}
}
}
#[cfg(test)]
mod tests {
use ruff_formatter::format;

610
crates/ruff_python_formatter/src/verbatim.rs Normal file
View file

@ -0,0 +1,610 @@
use std::borrow::Cow;
use std::iter::FusedIterator;
use ruff_formatter::write;
use ruff_python_ast::node::AnyNodeRef;
use ruff_python_ast::{Ranged, Stmt};
use ruff_python_trivia::lines_before;
use ruff_text_size::TextRange;
use crate::comments::format::{empty_lines, format_comment};
use crate::comments::{leading_comments, trailing_comments, SourceComment};
use crate::prelude::*;
use crate::statement::suite::SuiteChildStatement;
/// Disables formatting for all statements between the `first_suppressed` that has a leading `fmt: off` comment
/// and the first trailing or leading `fmt: on` comment. The statements are formatted as they appear in the source code.
///
/// Returns the last formatted statement.
///
/// ## Panics
/// If `first_suppressed` has no leading suppression comment.
#[cold]
pub(crate) fn write_suppressed_statements_starting_with_leading_comment<'a>(
// The first suppressed statement
first_suppressed: SuiteChildStatement<'a>,
statements: &mut std::slice::Iter<'a, Stmt>,
f: &mut PyFormatter,
) -> FormatResult<&'a Stmt> {
let comments = f.context().comments().clone();
let source = f.context().source();
let mut leading_comment_ranges =
CommentRangeIter::outside_suppression(comments.leading_comments(first_suppressed), source);
let before_format_off = leading_comment_ranges
.next()
.expect("Suppressed node to have leading comments");
let (formatted_comments, format_off_comment) = before_format_off.unwrap_suppression_starts();
// Format the leading comments before the fmt off
// ```python
// # leading comment that gets formatted
// # fmt: off
// statement
// ```
write!(
f,
[
leading_comments(formatted_comments),
// Format the off comment without adding any trailing new lines
format_comment(format_off_comment)
]
)?;
format_off_comment.mark_formatted();
// Now inside a suppressed range
write_suppressed_statements(
format_off_comment,
first_suppressed,
leading_comment_ranges.as_slice(),
statements,
f,
)
}
/// Disables formatting for all statements between the `last_formatted` and the first trailing or leading `fmt: on` comment.
/// The statements are formatted as they appear in the source code.
///
/// Returns the last formatted statement.
///
/// ## Panics
/// If `last_formatted` has no trailing suppression comment.
#[cold]
pub(crate) fn write_suppressed_statements_starting_with_trailing_comment<'a>(
last_formatted: SuiteChildStatement<'a>,
statements: &mut std::slice::Iter<'a, Stmt>,
f: &mut PyFormatter,
) -> FormatResult<&'a Stmt> {
let comments = f.context().comments().clone();
let source = f.context().source();
let trailing_node_comments = comments.trailing_comments(last_formatted);
let mut trailing_comment_ranges =
CommentRangeIter::outside_suppression(trailing_node_comments, source);
// Formatted comments gets formatted as part of the statement.
let (_, mut format_off_comment) = trailing_comment_ranges
.next()
.expect("Suppressed statement to have trailing comments")
.unwrap_suppression_starts();
let maybe_suppressed = trailing_comment_ranges.as_slice();
// Mark them as formatted so that calling the node's formatting doesn't format the comments.
for comment in maybe_suppressed {
comment.mark_formatted();
}
format_off_comment.mark_formatted();
// Format the leading comments, the node, and the trailing comments up to the `fmt: off` comment.
last_formatted.fmt(f)?;
format_off_comment.mark_unformatted();
TrailingFormatOffComment(format_off_comment).fmt(f)?;
for range in trailing_comment_ranges {
match range {
// A `fmt: off`..`fmt: on` sequence. Disable formatting for the in-between comments.
// ```python
// def test():
// pass
// # fmt: off
// # haha
// # fmt: on
// # fmt: off (maybe)
// ```
SuppressionComments::SuppressionEnds {
suppressed_comments: _,
format_on_comment,
formatted_comments,
format_off_comment: new_format_off_comment,
} => {
format_on_comment.mark_unformatted();
for comment in formatted_comments {
comment.mark_unformatted();
}
write!(
f,
[
verbatim_text(TextRange::new(
format_off_comment.end(),
format_on_comment.start(),
)),
trailing_comments(std::slice::from_ref(format_on_comment)),
trailing_comments(formatted_comments),
]
)?;
// `fmt: off`..`fmt:on`..`fmt:off` sequence
// ```python
// def test():
// pass
// # fmt: off
// # haha
// # fmt: on
// # fmt: off
// ```
if let Some(new_format_off_comment) = new_format_off_comment {
new_format_off_comment.mark_unformatted();
TrailingFormatOffComment(new_format_off_comment).fmt(f)?;
format_off_comment = new_format_off_comment;
} else {
// `fmt: off`..`fmt:on` sequence. The suppression ends here. Start formatting the nodes again.
return Ok(last_formatted.statement());
}
}
// All comments in this range are suppressed
SuppressionComments::Suppressed { comments: _ } => {}
// SAFETY: Unreachable because the function returns as soon as we reach the end of the suppressed range
SuppressionComments::SuppressionStarts { .. }
| SuppressionComments::Formatted { .. } => unreachable!(),
}
}
// The statement with the suppression comment isn't the last statement in the suite.
// Format the statements up to the first `fmt: on` comment (or end of the suite) as verbatim/suppressed.
// ```python
// a + b
// # fmt: off
//
// def a():
// pass
// ```
if let Some(first_suppressed) = statements.next() {
write_suppressed_statements(
format_off_comment,
SuiteChildStatement::Other(first_suppressed),
comments.leading_comments(first_suppressed),
statements,
f,
)
}
// The suppression comment is the block's last node. Format any trailing comments as suppressed
// ```python
// def test():
// pass
// # fmt: off
// # a trailing comment
// ```
else if let Some(last_comment) = trailing_node_comments.last() {
verbatim_text(TextRange::new(format_off_comment.end(), last_comment.end())).fmt(f)?;
Ok(last_formatted.statement())
}
// The suppression comment is the very last code in the block. There's nothing more to format.
// ```python
// def test():
// pass
// # fmt: off
// ```
else {
Ok(last_formatted.statement())
}
}
/// Formats the statements from `first_suppressed` until the suppression ends (by a `fmt: on` comment)
/// as they appear in the source code.
fn write_suppressed_statements<'a>(
// The `fmt: off` comment that starts the suppressed range. Can be a leading comment of `first_suppressed` or
// a trailing comment of the previous node.
format_off_comment: &SourceComment,
// The first suppressed statement
first_suppressed: SuiteChildStatement<'a>,
// The leading comments of `first_suppressed` that come after the `format_off_comment`
first_suppressed_leading_comments: &[SourceComment],
// The remaining statements
statements: &mut std::slice::Iter<'a, Stmt>,
f: &mut PyFormatter,
) -> FormatResult<&'a Stmt> {
let comments = f.context().comments().clone();
let source = f.context().source();
// TODO(micha) Fixup indent
let mut statement = first_suppressed;
let mut leading_node_comments = first_suppressed_leading_comments;
let mut format_off_comment = format_off_comment;
loop {
for range in CommentRangeIter::in_suppression(leading_node_comments, source) {
match range {
// All leading comments are suppressed
// ```python
// # suppressed comment
// statement
// ```
SuppressionComments::Suppressed { comments } => {
for comment in comments {
comment.mark_formatted();
}
}
// Node has a leading `fmt: on` comment and maybe another `fmt: off` comment
// ```python
// # suppressed comment (optional)
// # fmt: on
// # formatted comment (optional)
// # fmt: off (optional)
// statement
// ```
SuppressionComments::SuppressionEnds {
suppressed_comments,
format_on_comment,
formatted_comments,
format_off_comment: new_format_off_comment,
} => {
for comment in suppressed_comments {
comment.mark_formatted();
}
write!(
f,
[
verbatim_text(TextRange::new(
format_off_comment.end(),
format_on_comment.start(),
)),
leading_comments(std::slice::from_ref(format_on_comment)),
leading_comments(formatted_comments),
]
)?;
if let Some(new_format_off_comment) = new_format_off_comment {
format_off_comment = new_format_off_comment;
format_comment(format_off_comment).fmt(f)?;
format_off_comment.mark_formatted();
} else {
// Suppression ends here. Test if the node has a trailing suppression comment and, if so,
// recurse and format the trailing comments and the following statements as suppressed.
return if comments
.trailing_comments(statement)
.iter()
.any(|comment| comment.is_suppression_off_comment(source))
{
// Node has a trailing suppression comment, hell yeah, start all over again.
write_suppressed_statements_starting_with_trailing_comment(
statement, statements, f,
)
} else {
// Formats the trailing comments
statement.fmt(f)?;
Ok(statement.statement())
};
}
}
// Unreachable because the function exits as soon as it reaches the end of the suppression
// and it already starts in a suppressed range.
SuppressionComments::SuppressionStarts { .. } => unreachable!(),
SuppressionComments::Formatted { .. } => unreachable!(),
}
}
comments.mark_verbatim_node_comments_formatted(AnyNodeRef::from(statement));
for range in CommentRangeIter::in_suppression(comments.trailing_comments(statement), source)
{
match range {
// All leading comments are suppressed
// ```python
// statement
// # suppressed
// ```
SuppressionComments::Suppressed { comments } => {
for comment in comments {
comment.mark_formatted();
}
}
// Node has a trailing `fmt: on` comment and maybe another `fmt: off` comment
// ```python
// statement
// # suppressed comment (optional)
// # fmt: on
// # formatted comment (optional)
// # fmt: off (optional)
// ```
SuppressionComments::SuppressionEnds {
suppressed_comments,
format_on_comment,
formatted_comments,
format_off_comment: new_format_off_comment,
} => {
for comment in suppressed_comments {
comment.mark_formatted();
}
write!(
f,
[
verbatim_text(TextRange::new(
format_off_comment.end(),
format_on_comment.start()
)),
format_comment(format_on_comment),
hard_line_break(),
trailing_comments(formatted_comments),
]
)?;
format_on_comment.mark_formatted();
if let Some(new_format_off_comment) = new_format_off_comment {
format_off_comment = new_format_off_comment;
format_comment(format_off_comment).fmt(f)?;
format_off_comment.mark_formatted();
} else {
return Ok(statement.statement());
}
}
// Unreachable because the function exits as soon as it reaches the end of the suppression
// and it already starts in a suppressed range.
SuppressionComments::SuppressionStarts { .. } => unreachable!(),
SuppressionComments::Formatted { .. } => unreachable!(),
}
}
if let Some(next_statement) = statements.next() {
statement = SuiteChildStatement::Other(next_statement);
leading_node_comments = comments.leading_comments(next_statement);
} else {
let end = comments
.trailing_comments(statement)
.last()
.map_or(statement.end(), Ranged::end);
verbatim_text(TextRange::new(format_off_comment.end(), end)).fmt(f)?;
return Ok(statement.statement());
}
}
}
#[derive(Copy, Clone, Debug)]
enum InSuppression {
No,
Yes,
}
#[derive(Debug)]
enum SuppressionComments<'a> {
/// The first `fmt: off` comment.
SuppressionStarts {
/// The comments appearing before the `fmt: off` comment
formatted_comments: &'a [SourceComment],
format_off_comment: &'a SourceComment,
},
/// A `fmt: on` comment inside a suppressed range.
SuppressionEnds {
/// The comments before the `fmt: on` comment that should *not* be formatted.
suppressed_comments: &'a [SourceComment],
format_on_comment: &'a SourceComment,
/// The comments after the `fmt: on` comment (if any), that should be formatted.
formatted_comments: &'a [SourceComment],
/// Any following `fmt: off` comment if any.
/// * `None`: The suppression ends here (for good)
/// * `Some`: A `fmt: off`..`fmt: on` .. `fmt: off` sequence. The suppression continues after
/// the `fmt: off` comment.
format_off_comment: Option<&'a SourceComment>,
},
/// Comments that all fall into the suppressed range.
Suppressed { comments: &'a [SourceComment] },
/// Comments that all fall into the formatted range.
Formatted {
#[allow(unused)]
comments: &'a [SourceComment],
},
}
impl<'a> SuppressionComments<'a> {
fn unwrap_suppression_starts(&self) -> (&'a [SourceComment], &'a SourceComment) {
if let SuppressionComments::SuppressionStarts {
formatted_comments,
format_off_comment,
} = self
{
(formatted_comments, *format_off_comment)
} else {
panic!("Expected SuppressionStarts")
}
}
}
struct CommentRangeIter<'a> {
comments: &'a [SourceComment],
source: &'a str,
in_suppression: InSuppression,
}
impl<'a> CommentRangeIter<'a> {
fn in_suppression(comments: &'a [SourceComment], source: &'a str) -> Self {
Self {
comments,
in_suppression: InSuppression::Yes,
source,
}
}
fn outside_suppression(comments: &'a [SourceComment], source: &'a str) -> Self {
Self {
comments,
in_suppression: InSuppression::No,
source,
}
}
/// Returns a slice containing the remaining comments.
fn as_slice(&self) -> &'a [SourceComment] {
self.comments
}
}
impl<'a> Iterator for CommentRangeIter<'a> {
type Item = SuppressionComments<'a>;
fn next(&mut self) -> Option<Self::Item> {
if self.comments.is_empty() {
return None;
}
Some(match self.in_suppression {
// Inside of a suppressed range
InSuppression::Yes => {
if let Some(format_on_position) = self
.comments
.iter()
.position(|comment| comment.is_suppression_on_comment(self.source))
{
let (suppressed_comments, formatted) =
self.comments.split_at(format_on_position);
let (format_on_comment, rest) = formatted.split_first().unwrap();
let (formatted_comments, format_off_comment) =
if let Some(format_off_position) = rest
.iter()
.position(|comment| comment.is_suppression_off_comment(self.source))
{
let (formatted_comments, suppressed_comments) =
rest.split_at(format_off_position);
let (format_off_comment, rest) =
suppressed_comments.split_first().unwrap();
self.comments = rest;
(formatted_comments, Some(format_off_comment))
} else {
self.in_suppression = InSuppression::No;
self.comments = &[];
(rest, None)
};
SuppressionComments::SuppressionEnds {
suppressed_comments,
format_on_comment,
formatted_comments,
format_off_comment,
}
} else {
SuppressionComments::Suppressed {
comments: std::mem::take(&mut self.comments),
}
}
}
// Outside of a suppression
InSuppression::No => {
if let Some(format_off_position) = self
.comments
.iter()
.position(|comment| comment.is_suppression_off_comment(self.source))
{
self.in_suppression = InSuppression::Yes;
let (formatted_comments, suppressed) =
self.comments.split_at(format_off_position);
let format_off_comment = &suppressed[0];
self.comments = &suppressed[1..];
SuppressionComments::SuppressionStarts {
formatted_comments,
format_off_comment,
}
} else {
SuppressionComments::Formatted {
comments: std::mem::take(&mut self.comments),
}
}
}
})
}
}
impl FusedIterator for CommentRangeIter<'_> {}
struct TrailingFormatOffComment<'a>(&'a SourceComment);
impl Format<PyFormatContext<'_>> for TrailingFormatOffComment<'_> {
fn fmt(&self, f: &mut Formatter<PyFormatContext<'_>>) -> FormatResult<()> {
debug_assert!(self.0.is_unformatted());
let lines_before_comment = lines_before(self.0.start(), f.context().source());
write!(
f,
[empty_lines(lines_before_comment), format_comment(self.0)]
)?;
self.0.mark_formatted();
Ok(())
}
}
struct VerbatimText(TextRange);
fn verbatim_text<T>(item: T) -> VerbatimText
where
T: Ranged,
{
VerbatimText(item.range())
}
impl Format<PyFormatContext<'_>> for VerbatimText {
fn fmt(&self, f: &mut PyFormatter) -> FormatResult<()> {
f.write_element(FormatElement::Tag(Tag::StartVerbatim(
tag::VerbatimKind::Verbatim {
length: self.0.len(),
},
)))?;
match normalize_newlines(f.context().locator().slice(self.0), ['\r']) {
Cow::Borrowed(_) => {
write!(f, [source_text_slice(self.0, ContainsNewlines::Detect)])?;
}
Cow::Owned(cleaned) => {
write!(
f,
[
dynamic_text(&cleaned, Some(self.0.start())),
source_position(self.0.end())
]
)?;
}
}
f.write_element(FormatElement::Tag(Tag::EndVerbatim))
}
}