use std::borrow::Cow;
use std::iter::FusedIterator;
use std::slice::Iter;
use itertools::PeekingNext;
use ruff_formatter::{FormatError, write};
use ruff_python_ast::AnyNodeRef;
use ruff_python_ast::Stmt;
use ruff_python_ast::token::{Token as AstToken, TokenKind};
use ruff_python_trivia::lines_before;
use ruff_source_file::LineRanges;
use ruff_text_size::{Ranged, TextRange, TextSize};
use crate::comments::format::{empty_lines, format_comment};
use crate::comments::{SourceComment, leading_comments, trailing_comments};
use crate::prelude::*;
use crate::statement::clause::ClauseHeader;
use crate::statement::suite::SuiteChildStatement;
use crate::statement::trailing_semicolon;
pub(crate) fn starts_suppression(
leading_or_trailing_comments: &[SourceComment],
source: &str,
) -> bool {
let mut iter = CommentRangeIter::outside_suppression(leading_or_trailing_comments, source);
let _ = iter.by_ref().last();
matches!(iter.in_suppression, InSuppression::Yes)
}
pub(crate) fn ends_suppression(
leading_or_trailing_comments: &[SourceComment],
source: &str,
) -> bool {
let mut iter = CommentRangeIter::in_suppression(leading_or_trailing_comments, source);
let _ = iter.by_ref().last();
!matches!(iter.in_suppression, InSuppression::Yes)
}
#[cold]
pub(crate) fn write_suppressed_statements_starting_with_leading_comment<'a>(
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(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();
write!(
f,
[
leading_comments(formatted_comments),
format_comment(format_off_comment)
]
)?;
format_off_comment.mark_formatted();
write_suppressed_statements(
format_off_comment,
first_suppressed,
leading_comment_ranges.as_slice(),
statements,
f,
)
}
#[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 indentation = Indentation::from_stmt(last_formatted.statement(), source);
let trailing_node_comments = comments.trailing(last_formatted);
let mut trailing_comment_ranges =
CommentRangeIter::outside_suppression(trailing_node_comments, source);
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();
for comment in maybe_suppressed {
comment.mark_formatted();
}
format_off_comment.mark_formatted();
last_formatted.fmt(f)?;
format_off_comment.mark_unformatted();
TrailingFormatOffComment(format_off_comment).fmt(f)?;
for range in trailing_comment_ranges {
match range {
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,
[
FormatVerbatimStatementRange {
verbatim_range: TextRange::new(
format_off_comment.end(),
format_on_comment.start(),
),
indentation
},
trailing_comments(std::slice::from_ref(format_on_comment)),
trailing_comments(formatted_comments),
]
)?;
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 {
return Ok(last_formatted.statement());
}
}
SuppressionComments::Suppressed { comments: _ } => {}
SuppressionComments::SuppressionStarts { .. }
| SuppressionComments::Formatted { .. } => unreachable!(),
}
}
if let Some(first_suppressed) = statements.next() {
write_suppressed_statements(
format_off_comment,
SuiteChildStatement::Other(first_suppressed),
comments.leading(first_suppressed),
statements,
f,
)
}
else if let Some(last_comment) = trailing_node_comments.last() {
FormatVerbatimStatementRange {
verbatim_range: TextRange::new(format_off_comment.end(), last_comment.end()),
indentation,
}
.fmt(f)?;
Ok(last_formatted.statement())
}
else {
Ok(last_formatted.statement())
}
}
fn write_suppressed_statements<'a>(
format_off_comment: &SourceComment,
first_suppressed: SuiteChildStatement<'a>,
first_suppressed_leading_comments: &[SourceComment],
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 statement = first_suppressed;
let mut leading_node_comments = first_suppressed_leading_comments;
let mut format_off_comment = format_off_comment;
let indentation = Indentation::from_stmt(first_suppressed.statement(), source);
loop {
for range in CommentRangeIter::in_suppression(leading_node_comments, source) {
match range {
SuppressionComments::Suppressed { comments } => {
for comment in comments {
comment.mark_formatted();
}
}
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,
[
FormatVerbatimStatementRange {
verbatim_range: TextRange::new(
format_off_comment.end(),
format_on_comment.start(),
),
indentation
},
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 {
return if comments
.trailing(statement)
.iter()
.any(|comment| comment.is_suppression_off_comment(source))
{
write_suppressed_statements_starting_with_trailing_comment(
statement, statements, f,
)
} else {
statement.fmt(f)?;
Ok(statement.statement())
};
}
}
SuppressionComments::SuppressionStarts { .. } => unreachable!(),
SuppressionComments::Formatted { .. } => unreachable!(),
}
}
comments.mark_verbatim_node_comments_formatted(AnyNodeRef::from(statement));
for range in CommentRangeIter::in_suppression(comments.trailing(statement), source) {
match range {
SuppressionComments::Suppressed { comments } => {
for comment in comments {
comment.mark_formatted();
}
}
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,
[
FormatVerbatimStatementRange {
verbatim_range: TextRange::new(
format_off_comment.end(),
format_on_comment.start()
),
indentation
},
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());
}
}
SuppressionComments::SuppressionStarts { .. } => unreachable!(),
SuppressionComments::Formatted { .. } => unreachable!(),
}
}
if let Some(next_statement) = statements.next() {
statement = SuiteChildStatement::Other(next_statement);
leading_node_comments = comments.leading(next_statement);
} else {
let mut current = AnyNodeRef::from(statement.statement());
let end = loop {
if let Some(comment) = comments.trailing(current).last() {
break comment.end();
} else if let Some(child) = current.last_child_in_body() {
current = child;
} else {
break trailing_semicolon(current, source)
.map_or(statement.end(), TextRange::end);
}
};
FormatVerbatimStatementRange {
verbatim_range: TextRange::new(format_off_comment.end(), end),
indentation,
}
.fmt(f)?;
return Ok(statement.statement());
}
}
}
#[cold]
pub(crate) fn write_skipped_statements<'a>(
first_skipped: &'a Stmt,
statements: &mut std::slice::Iter<'a, Stmt>,
verbatim_range: TextRange,
f: &mut PyFormatter,
) -> FormatResult<&'a Stmt> {
let comments = f.context().comments().clone();
comments.mark_verbatim_node_comments_formatted(first_skipped.into());
let mut preceding = first_skipped;
while let Some(prec) = statements.peeking_next(|next| next.end() <= verbatim_range.end()) {
comments.mark_verbatim_node_comments_formatted(prec.into());
preceding = prec;
}
let first_leading = comments.leading(first_skipped);
let preceding_trailing = comments.trailing(preceding);
write!(
f,
[
leading_comments(first_leading),
source_position(verbatim_range.start()),
verbatim_text(verbatim_range),
source_position(verbatim_range.end()),
trailing_comments(preceding_trailing)
]
)?;
Ok(preceding)
}
#[derive(Copy, Clone, Debug)]
enum InSuppression {
No,
Yes,
}
#[derive(Debug)]
enum SuppressionComments<'a> {
SuppressionStarts {
formatted_comments: &'a [SourceComment],
format_off_comment: &'a SourceComment,
},
SuppressionEnds {
suppressed_comments: &'a [SourceComment],
format_on_comment: &'a SourceComment,
formatted_comments: &'a [SourceComment],
format_off_comment: Option<&'a SourceComment>,
},
Suppressed { comments: &'a [SourceComment] },
Formatted {
#[expect(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,
}
}
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 {
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),
}
}
}
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(())
}
}
#[derive(Copy, Clone)]
struct Indentation(u32);
impl Indentation {
fn from_stmt(stmt: &Stmt, source: &str) -> Indentation {
let line_start = source.line_start(stmt.start());
let mut indentation = 0u32;
for c in source[TextRange::new(line_start, stmt.start())].chars() {
if is_indent_whitespace(c) {
indentation += 1;
} else {
break;
}
}
Indentation(indentation)
}
fn trim_indent(self, ranged: impl Ranged, source: &str) -> TextRange {
let range = ranged.range();
let mut start_offset = TextSize::default();
for c in source[range].chars().take(self.0 as usize) {
if is_indent_whitespace(c) {
start_offset += TextSize::new(1);
} else {
break;
}
}
TextRange::new(range.start() + start_offset, range.end())
}
}
const fn is_indent_whitespace(c: char) -> bool {
matches!(c, ' ' | '\t')
}
struct FormatVerbatimStatementRange {
verbatim_range: TextRange,
indentation: Indentation,
}
impl Format<PyFormatContext<'_>> for FormatVerbatimStatementRange {
fn fmt(&self, f: &mut Formatter<PyFormatContext<'_>>) -> FormatResult<()> {
let logical_lines = LogicalLinesIter::new(
f.context().tokens().in_range(self.verbatim_range).iter(),
self.verbatim_range,
);
let mut first = true;
for logical_line in logical_lines {
let logical_line = logical_line?;
let trimmed_line_range = self
.indentation
.trim_indent(&logical_line, f.context().source());
if trimmed_line_range.is_empty() {
if logical_line.has_trailing_newline {
if first {
hard_line_break().fmt(f)?;
} else {
empty_line().fmt(f)?;
}
}
} else {
write!(
f,
[
source_position(trimmed_line_range.start()),
verbatim_text(trimmed_line_range),
source_position(trimmed_line_range.end())
]
)?;
if logical_line.has_trailing_newline {
hard_line_break().fmt(f)?;
}
}
first = false;
}
Ok(())
}
}
struct LogicalLinesIter<'a> {
tokens: Iter<'a, AstToken>,
last_line_end: TextSize,
content_end: TextSize,
}
impl<'a> LogicalLinesIter<'a> {
fn new(tokens: Iter<'a, AstToken>, verbatim_range: TextRange) -> Self {
Self {
tokens,
last_line_end: verbatim_range.start(),
content_end: verbatim_range.end(),
}
}
}
impl Iterator for LogicalLinesIter<'_> {
type Item = FormatResult<LogicalLine>;
fn next(&mut self) -> Option<Self::Item> {
let mut parens = 0u32;
let (content_end, full_end) = loop {
match self.tokens.next() {
Some(token) if token.kind() == TokenKind::Unknown => {
return Some(Err(FormatError::syntax_error(
"Unexpected token when lexing verbatim statement range.",
)));
}
Some(token) => match token.kind() {
TokenKind::Newline => break (token.start(), token.end()),
TokenKind::NonLogicalNewline if parens == 0 => {
break (token.start(), token.end());
}
TokenKind::Lbrace | TokenKind::Lpar | TokenKind::Lsqb => {
parens = parens.saturating_add(1);
}
TokenKind::Rbrace | TokenKind::Rpar | TokenKind::Rsqb => {
parens = parens.saturating_sub(1);
}
_ => {}
},
None => {
return if self.last_line_end < self.content_end {
let content_start = self.last_line_end;
self.last_line_end = self.content_end;
Some(Ok(LogicalLine {
content_range: TextRange::new(content_start, self.content_end),
has_trailing_newline: false,
}))
} else {
None
};
}
}
};
let line_start = self.last_line_end;
self.last_line_end = full_end;
Some(Ok(LogicalLine {
content_range: TextRange::new(line_start, content_end),
has_trailing_newline: true,
}))
}
}
impl FusedIterator for LogicalLinesIter<'_> {}
struct LogicalLine {
content_range: TextRange,
has_trailing_newline: bool,
}
impl Ranged for LogicalLine {
fn range(&self) -> TextRange {
self.content_range
}
}
pub(crate) struct VerbatimText {
verbatim_range: TextRange,
}
pub(crate) fn verbatim_text<T>(item: T) -> VerbatimText
where
T: Ranged,
{
VerbatimText {
verbatim_range: 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.verbatim_range.len(),
},
)));
match normalize_newlines(&f.context().source()[self.verbatim_range], ['\r']) {
Cow::Borrowed(_) => {
write!(f, [source_text_slice(self.verbatim_range)])?;
}
Cow::Owned(cleaned) => {
text(&cleaned).fmt(f)?;
}
}
f.write_element(FormatElement::Tag(Tag::EndVerbatim));
Ok(())
}
}
#[cold]
pub(crate) fn write_suppressed_clause_header(
header: ClauseHeader,
f: &mut PyFormatter,
) -> FormatResult<()> {
let range = header.range(f.context().source())?;
write!(
f,
[
source_position(range.start()),
verbatim_text(range),
source_position(range.end())
]
)?;
let comments = f.context().comments();
header.visit(&mut |child| {
for comment in comments.leading_trailing(child) {
comment.mark_formatted();
}
comments.mark_verbatim_node_comments_formatted(child);
});
Ok(())
}