mirror of
https://github.com/astral-sh/ruff.git
synced 2025-09-30 22:01:47 +00:00
9a817a2922
When there is a function or class definition at the end of a suite
followed by the beginning of an alternative block, we have to insert a
single empty line between them.
In the if-else-statement example below, we insert an empty line after
the `foo` in the if-block, but none after the else-block `foo`, since in
the latter case the enclosing suite already adds empty lines.
```python
if sys.version_info >= (3, 10):
def foo():
return "new"
else:
def foo():
return "old"
class Bar:
pass
```
To do so, we track whether the current suite is the last one in the
current statement with a new option on the suite kind.
Fixes #12199
---------
Co-authored-by: Micha Reiser <micha@reiser.io>
1045 lines
36 KiB
Rust
1045 lines
36 KiB
Rust
use ruff_formatter::{
|
|
write, FormatContext, FormatOwnedWithRule, FormatRefWithRule, FormatRuleWithOptions,
|
|
};
|
|
use ruff_python_ast::helpers::is_compound_statement;
|
|
use ruff_python_ast::{self as ast, Expr, PySourceType, Stmt, Suite};
|
|
use ruff_python_ast::{AnyNodeRef, StmtExpr};
|
|
use ruff_python_trivia::{lines_after, lines_after_ignoring_end_of_line_trivia, lines_before};
|
|
use ruff_text_size::{Ranged, TextRange};
|
|
|
|
use crate::comments::{
|
|
leading_comments, trailing_comments, Comments, LeadingDanglingTrailingComments,
|
|
};
|
|
use crate::context::{NodeLevel, TopLevelStatementPosition, WithIndentLevel, WithNodeLevel};
|
|
use crate::expression::expr_string_literal::ExprStringLiteralKind;
|
|
use crate::prelude::*;
|
|
use crate::statement::stmt_expr::FormatStmtExpr;
|
|
use crate::verbatim::{
|
|
suppressed_node, 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, PartialEq, Eq)]
|
|
pub enum SuiteKind {
|
|
/// Statements at the module level / top level
|
|
TopLevel,
|
|
|
|
/// Statements in a function body.
|
|
Function,
|
|
|
|
/// Statements in a class body.
|
|
Class,
|
|
|
|
/// Statements in any other body (e.g., `if` or `while`).
|
|
Other {
|
|
/// Whether this suite is the last suite in the current statement.
|
|
///
|
|
/// Below, `last_suite_in_statement` is `false` for the suite containing `foo10` and `foo12`
|
|
/// and `true` for the suite containing `bar`.
|
|
/// ```python
|
|
/// if sys.version_info >= (3, 10):
|
|
/// def foo10():
|
|
/// return "new"
|
|
/// elif sys.version_info >= (3, 12):
|
|
/// def foo12():
|
|
/// return "new"
|
|
/// else:
|
|
/// def bar():
|
|
/// return "old"
|
|
/// ```
|
|
///
|
|
/// When this value is true, we don't insert trailing empty lines since the containing suite
|
|
/// will do that.
|
|
last_suite_in_statement: bool,
|
|
},
|
|
}
|
|
|
|
impl Default for SuiteKind {
|
|
fn default() -> Self {
|
|
Self::Other {
|
|
// For stability, we can't insert an empty line if we don't know if the outer suite
|
|
// also does.
|
|
last_suite_in_statement: true,
|
|
}
|
|
}
|
|
}
|
|
|
|
impl SuiteKind {
|
|
/// See [`SuiteKind::Other`].
|
|
pub fn other(last_suite_in_statement: bool) -> Self {
|
|
Self::Other {
|
|
last_suite_in_statement,
|
|
}
|
|
}
|
|
|
|
pub fn last_suite_in_statement(self) -> bool {
|
|
match self {
|
|
Self::Other {
|
|
last_suite_in_statement,
|
|
} => last_suite_in_statement,
|
|
_ => true,
|
|
}
|
|
}
|
|
}
|
|
|
|
#[derive(Debug, Default)]
|
|
pub struct FormatSuite {
|
|
kind: SuiteKind,
|
|
}
|
|
|
|
impl FormatRule<Suite, PyFormatContext<'_>> for FormatSuite {
|
|
fn fmt(&self, statements: &Suite, f: &mut PyFormatter) -> FormatResult<()> {
|
|
let mut iter = statements.iter();
|
|
let Some(first) = iter.next() else {
|
|
return Ok(());
|
|
};
|
|
|
|
let node_level = match self.kind {
|
|
SuiteKind::TopLevel => NodeLevel::TopLevel(
|
|
iter.clone()
|
|
.next()
|
|
.map_or(TopLevelStatementPosition::Last, |_| {
|
|
TopLevelStatementPosition::Other
|
|
}),
|
|
),
|
|
SuiteKind::Function | SuiteKind::Class | SuiteKind::Other { .. } => {
|
|
NodeLevel::CompoundStatement
|
|
}
|
|
};
|
|
|
|
let comments = f.context().comments().clone();
|
|
let source = f.context().source();
|
|
let source_type = f.options().source_type();
|
|
|
|
let f = WithNodeLevel::new(node_level, f);
|
|
let f = &mut WithIndentLevel::new(f.context().indent_level().increment(), f);
|
|
|
|
// 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(first)
|
|
&& !source_type.is_stub()
|
|
{
|
|
// 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 | SuiteKind::Class | SuiteKind::TopLevel => {
|
|
if let Some(docstring) =
|
|
DocstringStmt::try_from_statement(first, self.kind, source_type)
|
|
{
|
|
SuiteChildStatement::Docstring(docstring)
|
|
} else {
|
|
SuiteChildStatement::Other(first)
|
|
}
|
|
}
|
|
};
|
|
|
|
let first_comments = comments.leading_dangling_trailing(first);
|
|
|
|
let (mut preceding, mut empty_line_after_docstring) = if first_comments
|
|
.leading
|
|
.iter()
|
|
.any(|comment| comment.is_suppression_off_comment(source))
|
|
{
|
|
(
|
|
write_suppressed_statements_starting_with_leading_comment(first, &mut iter, f)?,
|
|
false,
|
|
)
|
|
} else if first_comments
|
|
.trailing
|
|
.iter()
|
|
.any(|comment| comment.is_suppression_off_comment(source))
|
|
{
|
|
(
|
|
write_suppressed_statements_starting_with_trailing_comment(first, &mut iter, f)?,
|
|
false,
|
|
)
|
|
} else {
|
|
first.fmt(f)?;
|
|
|
|
#[allow(clippy::if_same_then_else)]
|
|
let empty_line_after_docstring = if matches!(first, SuiteChildStatement::Docstring(_))
|
|
&& self.kind == SuiteKind::Class
|
|
{
|
|
true
|
|
} else {
|
|
// Insert a newline after a module level docstring, but treat
|
|
// it as a docstring otherwise. See: https://github.com/psf/black/pull/3932.
|
|
self.kind == SuiteKind::TopLevel
|
|
&& DocstringStmt::try_from_statement(first.statement(), self.kind, source_type)
|
|
.is_some()
|
|
};
|
|
|
|
(first.statement(), empty_line_after_docstring)
|
|
};
|
|
|
|
let mut preceding_comments = comments.leading_dangling_trailing(preceding);
|
|
|
|
while let Some(following) = iter.next() {
|
|
if self.kind == SuiteKind::TopLevel && iter.clone().next().is_none() {
|
|
f.context_mut()
|
|
.set_node_level(NodeLevel::TopLevel(TopLevelStatementPosition::Last));
|
|
}
|
|
|
|
let following_comments = comments.leading_dangling_trailing(following);
|
|
|
|
let needs_empty_lines = if is_class_or_function_definition(following) {
|
|
// Here we insert empty lines even if the preceding has a trailing own line comment
|
|
true
|
|
} else {
|
|
trailing_function_or_class_def(Some(preceding), &comments).is_some()
|
|
};
|
|
|
|
// Add empty lines before and after a function or class definition. If the preceding
|
|
// node is a function or class, and contains trailing comments, then the statement
|
|
// itself will add the requisite empty lines when formatting its comments.
|
|
if needs_empty_lines {
|
|
if source_type.is_stub() {
|
|
stub_file_empty_lines(
|
|
self.kind,
|
|
preceding,
|
|
following,
|
|
&preceding_comments,
|
|
&following_comments,
|
|
f,
|
|
)?;
|
|
} else {
|
|
// Preserve empty lines after a stub implementation but don't insert a new one if there isn't any present in the source.
|
|
// This is useful when having multiple function overloads that should be grouped to getter by omitting new lines between them.
|
|
let is_preceding_stub_function_without_empty_line = following
|
|
.is_function_def_stmt()
|
|
&& preceding
|
|
.as_function_def_stmt()
|
|
.is_some_and(|preceding_stub| {
|
|
contains_only_an_ellipsis(
|
|
&preceding_stub.body,
|
|
f.context().comments(),
|
|
) && lines_after_ignoring_end_of_line_trivia(
|
|
preceding_stub.end(),
|
|
f.context().source(),
|
|
) < 2
|
|
})
|
|
&& !preceding_comments.has_trailing_own_line();
|
|
|
|
if !is_preceding_stub_function_without_empty_line {
|
|
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) || following_comments.has_leading())
|
|
{
|
|
// Enforce _at least_ one empty line after an import statement (but allow up to
|
|
// two at the top-level). In this context, "after an import statement" means that
|
|
// that the previous node is an import, and the following node is an import _or_ has
|
|
// a leading comment.
|
|
match self.kind {
|
|
SuiteKind::TopLevel => {
|
|
let end = if let Some(last_trailing) = preceding_comments.trailing.last() {
|
|
last_trailing.end()
|
|
} else {
|
|
preceding.end()
|
|
};
|
|
match lines_after(end, source) {
|
|
0..=2 => empty_line().fmt(f)?,
|
|
_ => match source_type {
|
|
PySourceType::Stub => {
|
|
empty_line().fmt(f)?;
|
|
}
|
|
PySourceType::Python | PySourceType::Ipynb => {
|
|
write!(f, [empty_line(), empty_line()])?;
|
|
}
|
|
},
|
|
}
|
|
}
|
|
SuiteKind::Function | SuiteKind::Class | SuiteKind::Other { .. } => {
|
|
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) = following_comments.leading.first() {
|
|
first_leading.start()
|
|
} else {
|
|
following.start()
|
|
};
|
|
|
|
match lines_before(start, source) {
|
|
0 | 1 => hard_line_break().fmt(f)?,
|
|
2 => empty_line().fmt(f)?,
|
|
_ => match self.kind {
|
|
SuiteKind::TopLevel => match source_type {
|
|
PySourceType::Stub => {
|
|
empty_line().fmt(f)?;
|
|
}
|
|
PySourceType::Python | PySourceType::Ipynb => {
|
|
write!(f, [empty_line(), empty_line()])?;
|
|
}
|
|
},
|
|
SuiteKind::Function | SuiteKind::Class | SuiteKind::Other { .. } => {
|
|
empty_line().fmt(f)?;
|
|
}
|
|
},
|
|
}
|
|
} else if empty_line_after_docstring {
|
|
// Enforce an empty line after a class docstring, e.g., these are both stable
|
|
// formatting:
|
|
// ```python
|
|
// class Test:
|
|
// """Docstring"""
|
|
//
|
|
// ...
|
|
//
|
|
//
|
|
// class Test:
|
|
//
|
|
// """Docstring"""
|
|
//
|
|
// ...
|
|
// ```
|
|
empty_line().fmt(f)?;
|
|
} 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
|
|
|
|
// It's necessary to skip any trailing line comment because our parser doesn't
|
|
// include trailing comments in the node's range:
|
|
// ```python
|
|
// a # The range of `a` ends right before this comment
|
|
//
|
|
// b
|
|
// ```
|
|
let end = preceding_comments
|
|
.trailing
|
|
.last()
|
|
.map_or(preceding.end(), |comment| comment.slice().end());
|
|
|
|
match node_level {
|
|
NodeLevel::TopLevel(_) => match lines_after(end, source) {
|
|
0 | 1 => hard_line_break().fmt(f)?,
|
|
2 => empty_line().fmt(f)?,
|
|
_ => match source_type {
|
|
PySourceType::Stub => {
|
|
empty_line().fmt(f)?;
|
|
}
|
|
PySourceType::Python | PySourceType::Ipynb => {
|
|
write!(f, [empty_line(), empty_line()])?;
|
|
}
|
|
},
|
|
},
|
|
NodeLevel::CompoundStatement => match lines_after(end, source) {
|
|
0 | 1 => hard_line_break().fmt(f)?,
|
|
_ => empty_line().fmt(f)?,
|
|
},
|
|
NodeLevel::Expression(_) | NodeLevel::ParenthesizedExpression => {
|
|
hard_line_break().fmt(f)?;
|
|
}
|
|
}
|
|
}
|
|
|
|
if following_comments
|
|
.leading
|
|
.iter()
|
|
.any(|comment| comment.is_suppression_off_comment(source))
|
|
{
|
|
preceding = write_suppressed_statements_starting_with_leading_comment(
|
|
SuiteChildStatement::Other(following),
|
|
&mut iter,
|
|
f,
|
|
)?;
|
|
preceding_comments = comments.leading_dangling_trailing(preceding);
|
|
} else if following_comments
|
|
.trailing
|
|
.iter()
|
|
.any(|comment| comment.is_suppression_off_comment(source))
|
|
{
|
|
preceding = write_suppressed_statements_starting_with_trailing_comment(
|
|
SuiteChildStatement::Other(following),
|
|
&mut iter,
|
|
f,
|
|
)?;
|
|
preceding_comments = comments.leading_dangling_trailing(preceding);
|
|
} else {
|
|
following.format().fmt(f)?;
|
|
preceding = following;
|
|
preceding_comments = following_comments;
|
|
}
|
|
|
|
empty_line_after_docstring = false;
|
|
}
|
|
|
|
self.between_alternative_blocks_empty_line(statements, &comments, f)?;
|
|
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
impl FormatSuite {
|
|
/// Add an empty line between a function or class and an alternative body.
|
|
///
|
|
/// We only insert an empty if we're between suites in a multi-suite statement. In the
|
|
/// if-else-statement below, we insert an empty line after the `foo` in the if-block, but none
|
|
/// after the else-block `foo`, since in the latter case the enclosing suite already adds
|
|
/// empty lines.
|
|
///
|
|
/// ```python
|
|
/// if sys.version_info >= (3, 10):
|
|
/// def foo():
|
|
/// return "new"
|
|
/// else:
|
|
/// def foo():
|
|
/// return "old"
|
|
/// class Bar:
|
|
/// pass
|
|
/// ```
|
|
fn between_alternative_blocks_empty_line(
|
|
&self,
|
|
statements: &Suite,
|
|
comments: &Comments,
|
|
f: &mut PyFormatter,
|
|
) -> FormatResult<()> {
|
|
if self.kind.last_suite_in_statement() {
|
|
// If we're at the end of the current statement, the outer suite will insert one or
|
|
// two empty lines already.
|
|
return Ok(());
|
|
}
|
|
|
|
let Some(last_def_or_class) = trailing_function_or_class_def(statements.last(), comments)
|
|
else {
|
|
// An empty line is only inserted for function and class definitions.
|
|
return Ok(());
|
|
};
|
|
|
|
// Skip the last trailing own line comment of the suite, if any, otherwise we count
|
|
// the lines wrongly by stopping at that comment.
|
|
let node_with_last_trailing_comment = std::iter::successors(
|
|
statements.last().map(AnyNodeRef::from),
|
|
AnyNodeRef::last_child_in_body,
|
|
)
|
|
.find(|last_child| comments.has_trailing_own_line(*last_child));
|
|
|
|
let end_of_def_or_class = node_with_last_trailing_comment
|
|
.and_then(|child| comments.trailing(child).last().map(Ranged::end))
|
|
.unwrap_or(last_def_or_class.end());
|
|
let existing_newlines =
|
|
lines_after_ignoring_end_of_line_trivia(end_of_def_or_class, f.context().source());
|
|
if existing_newlines < 2 {
|
|
if f.context().is_preview() {
|
|
empty_line().fmt(f)?;
|
|
} else {
|
|
if last_def_or_class.is_stmt_class_def() && f.options().source_type().is_stub() {
|
|
empty_line().fmt(f)?;
|
|
}
|
|
}
|
|
}
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
/// Find nested class or function definitions that need an empty line after them.
|
|
///
|
|
/// ```python
|
|
/// def f():
|
|
/// if True:
|
|
///
|
|
/// def double(s):
|
|
/// return s + s
|
|
///
|
|
/// print("below function")
|
|
/// ```
|
|
fn trailing_function_or_class_def<'a>(
|
|
preceding: Option<&'a Stmt>,
|
|
comments: &Comments,
|
|
) -> Option<AnyNodeRef<'a>> {
|
|
std::iter::successors(
|
|
preceding.map(AnyNodeRef::from),
|
|
AnyNodeRef::last_child_in_body,
|
|
)
|
|
.take_while(|last_child|
|
|
// If there is a comment between preceding and following the empty lines were
|
|
// inserted before the comment by preceding and there are no extra empty lines
|
|
// after the comment.
|
|
// ```python
|
|
// class Test:
|
|
// def a(self):
|
|
// pass
|
|
// # trailing comment
|
|
//
|
|
//
|
|
// # two lines before, one line after
|
|
//
|
|
// c = 30
|
|
// ````
|
|
// This also includes nested class/function definitions, so we stop recursing
|
|
// once we see a node with a trailing own line comment:
|
|
// ```python
|
|
// def f():
|
|
// if True:
|
|
//
|
|
// def double(s):
|
|
// return s + s
|
|
//
|
|
// # nested trailing own line comment
|
|
// print("below function with trailing own line comment")
|
|
// ```
|
|
!comments.has_trailing_own_line(*last_child))
|
|
.find(|last_child| {
|
|
matches!(
|
|
last_child,
|
|
AnyNodeRef::StmtFunctionDef(_) | AnyNodeRef::StmtClassDef(_)
|
|
)
|
|
})
|
|
}
|
|
|
|
/// Stub files have bespoke rules for empty lines.
|
|
///
|
|
/// These rules are ported from black (preview mode at time of writing) using the stubs test case:
|
|
/// <https://github.com/psf/black/blob/c160e4b7ce30c661ac4f2dfa5038becf1b8c5c33/src/black/lines.py#L576-L744>
|
|
fn stub_file_empty_lines(
|
|
kind: SuiteKind,
|
|
preceding: &Stmt,
|
|
following: &Stmt,
|
|
preceding_comments: &LeadingDanglingTrailingComments,
|
|
following_comments: &LeadingDanglingTrailingComments,
|
|
f: &mut PyFormatter,
|
|
) -> FormatResult<()> {
|
|
let source = f.context().source();
|
|
// Preserve the empty line if the definitions are separated by a comment
|
|
let empty_line_condition = preceding_comments.has_trailing()
|
|
|| following_comments.has_leading()
|
|
|| !stub_suite_can_omit_empty_line(preceding, following, f);
|
|
let require_empty_line = should_insert_blank_line_after_class_in_stub_file(
|
|
preceding.into(),
|
|
Some(following.into()),
|
|
f.context(),
|
|
);
|
|
match kind {
|
|
SuiteKind::TopLevel => {
|
|
if empty_line_condition || require_empty_line {
|
|
empty_line().fmt(f)
|
|
} else {
|
|
hard_line_break().fmt(f)
|
|
}
|
|
}
|
|
SuiteKind::Class | SuiteKind::Other { .. } | SuiteKind::Function => {
|
|
if (empty_line_condition
|
|
&& lines_after_ignoring_end_of_line_trivia(preceding.end(), source) > 1)
|
|
|| require_empty_line
|
|
{
|
|
empty_line().fmt(f)
|
|
} else {
|
|
hard_line_break().fmt(f)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Checks if an empty line should be inserted after a class definition.
|
|
///
|
|
/// This is only valid if the [`blank_line_after_nested_stub_class`](https://github.com/astral-sh/ruff/issues/8891)
|
|
/// preview rule is enabled and the source to be formatted is a stub file.
|
|
///
|
|
/// If `following` is `None`, then the preceding node is the last one in a suite. The
|
|
/// caller needs to make sure that the suite which the preceding node is part of is
|
|
/// followed by an alternate branch and shouldn't be a top-level suite.
|
|
pub(crate) fn should_insert_blank_line_after_class_in_stub_file(
|
|
preceding: AnyNodeRef<'_>,
|
|
following: Option<AnyNodeRef<'_>>,
|
|
context: &PyFormatContext,
|
|
) -> bool {
|
|
if !context.options().source_type().is_stub() {
|
|
return false;
|
|
}
|
|
|
|
let Some(following) = following else {
|
|
// We handle newlines at the end of a suite in `between_alternative_blocks_empty_line`.
|
|
return false;
|
|
};
|
|
|
|
let comments = context.comments();
|
|
match preceding.as_stmt_class_def() {
|
|
Some(class) if contains_only_an_ellipsis(&class.body, comments) => {
|
|
// If the preceding class has decorators, then we need to add an empty
|
|
// line even if it only contains ellipsis.
|
|
//
|
|
// ```python
|
|
// class Top:
|
|
// @decorator
|
|
// class Nested1: ...
|
|
// foo = 1
|
|
// ```
|
|
let preceding_has_decorators = !class.decorator_list.is_empty();
|
|
|
|
// If the following statement is a class definition, then an empty line
|
|
// should be inserted if it (1) doesn't just contain ellipsis, or (2) has decorators.
|
|
//
|
|
// ```python
|
|
// class Top:
|
|
// class Nested1: ...
|
|
// class Nested2:
|
|
// pass
|
|
//
|
|
// class Top:
|
|
// class Nested1: ...
|
|
// @decorator
|
|
// class Nested2: ...
|
|
// ```
|
|
//
|
|
// Both of the above examples should add a blank line in between.
|
|
let following_is_class_without_only_ellipsis_or_has_decorators =
|
|
following.as_stmt_class_def().is_some_and(|following| {
|
|
!contains_only_an_ellipsis(&following.body, comments)
|
|
|| !following.decorator_list.is_empty()
|
|
});
|
|
|
|
preceding_has_decorators
|
|
|| following_is_class_without_only_ellipsis_or_has_decorators
|
|
|| following.is_stmt_function_def()
|
|
}
|
|
Some(_) => {
|
|
// Preceding statement is a class definition whose body isn't only an ellipsis.
|
|
// Here, we should only add a blank line if the class doesn't have a trailing
|
|
// own line comment as that's handled by the class formatting itself.
|
|
!comments.has_trailing_own_line(preceding)
|
|
}
|
|
None => {
|
|
// If preceding isn't a class definition, let's check if the last statement
|
|
// in the body, going all the way down, is a class definition.
|
|
//
|
|
// ```python
|
|
// if foo:
|
|
// if bar:
|
|
// class Nested:
|
|
// pass
|
|
// if other:
|
|
// pass
|
|
// ```
|
|
//
|
|
// But, if it contained a trailing own line comment, then it's handled
|
|
// by the class formatting itself.
|
|
//
|
|
// ```python
|
|
// if foo:
|
|
// if bar:
|
|
// class Nested:
|
|
// pass
|
|
// # comment
|
|
// if other:
|
|
// pass
|
|
// ```
|
|
std::iter::successors(
|
|
preceding.last_child_in_body(),
|
|
AnyNodeRef::last_child_in_body,
|
|
)
|
|
.take_while(|last_child| !comments.has_trailing_own_line(*last_child))
|
|
.any(|last_child| last_child.is_stmt_class_def())
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Only a function to compute it lazily
|
|
fn stub_suite_can_omit_empty_line(preceding: &Stmt, following: &Stmt, f: &PyFormatter) -> bool {
|
|
// Two subsequent class definitions that both have an ellipsis only body
|
|
// ```python
|
|
// class A: ...
|
|
// class B: ...
|
|
//
|
|
// @decorator
|
|
// class C: ...
|
|
// ```
|
|
let class_sequences_with_ellipsis_only = preceding
|
|
.as_class_def_stmt()
|
|
.is_some_and(|class| contains_only_an_ellipsis(&class.body, f.context().comments()))
|
|
&& following.as_class_def_stmt().is_some_and(|class| {
|
|
contains_only_an_ellipsis(&class.body, f.context().comments())
|
|
&& class.decorator_list.is_empty()
|
|
});
|
|
|
|
// Black for some reasons accepts decorators in place of empty lines
|
|
// ```python
|
|
// def _count1(): ...
|
|
// @final
|
|
// class LockType1: ...
|
|
//
|
|
// def _count2(): ...
|
|
//
|
|
// class LockType2: ...
|
|
// ```
|
|
let class_decorator_instead_of_empty_line = preceding.is_function_def_stmt()
|
|
&& following
|
|
.as_class_def_stmt()
|
|
.is_some_and(|class| !class.decorator_list.is_empty());
|
|
|
|
// A function definition following a stub function definition
|
|
// ```python
|
|
// def test(): ...
|
|
// def b(): a
|
|
// ```
|
|
let function_with_ellipsis = preceding
|
|
.as_function_def_stmt()
|
|
.is_some_and(|function| contains_only_an_ellipsis(&function.body, f.context().comments()))
|
|
&& following.is_function_def_stmt();
|
|
|
|
class_sequences_with_ellipsis_only
|
|
|| class_decorator_instead_of_empty_line
|
|
|| function_with_ellipsis
|
|
}
|
|
|
|
/// Returns `true` if a function or class body contains only an ellipsis with no comments.
|
|
pub(crate) fn contains_only_an_ellipsis(body: &[Stmt], comments: &Comments) -> bool {
|
|
match body {
|
|
[Stmt::Expr(ast::StmtExpr { value, .. })] => {
|
|
let [node] = body else {
|
|
return false;
|
|
};
|
|
value.is_ellipsis_literal_expr()
|
|
&& !comments.has_leading(node)
|
|
&& !comments.has_trailing_own_line(node)
|
|
}
|
|
_ => false,
|
|
}
|
|
}
|
|
|
|
/// Returns `true` if a [`Stmt`] is a class or function definition.
|
|
const fn is_class_or_function_definition(stmt: &Stmt) -> bool {
|
|
matches!(stmt, Stmt::FunctionDef(_) | Stmt::ClassDef(_))
|
|
}
|
|
|
|
/// Returns `true` if a [`Stmt`] is an import.
|
|
const fn is_import_definition(stmt: &Stmt) -> bool {
|
|
matches!(stmt, Stmt::Import(_) | Stmt::ImportFrom(_))
|
|
}
|
|
|
|
impl FormatRuleWithOptions<Suite, PyFormatContext<'_>> for FormatSuite {
|
|
type Options = SuiteKind;
|
|
|
|
fn with_options(mut self, options: Self::Options) -> Self {
|
|
self.kind = options;
|
|
self
|
|
}
|
|
}
|
|
|
|
impl<'ast> AsFormat<PyFormatContext<'ast>> for Suite {
|
|
type Format<'a> = FormatRefWithRule<'a, Suite, FormatSuite, PyFormatContext<'ast>>;
|
|
|
|
fn format(&self) -> Self::Format<'_> {
|
|
FormatRefWithRule::new(self, FormatSuite::default())
|
|
}
|
|
}
|
|
|
|
impl<'ast> IntoFormat<PyFormatContext<'ast>> for Suite {
|
|
type Format = FormatOwnedWithRule<Suite, FormatSuite, PyFormatContext<'ast>>;
|
|
|
|
fn into_format(self) -> Self::Format {
|
|
FormatOwnedWithRule::new(self, FormatSuite::default())
|
|
}
|
|
}
|
|
|
|
/// A statement representing a docstring.
|
|
#[derive(Copy, Clone, Debug)]
|
|
pub(crate) struct DocstringStmt<'a> {
|
|
/// The [`Stmt::Expr`]
|
|
docstring: &'a Stmt,
|
|
/// The parent suite kind
|
|
suite_kind: SuiteKind,
|
|
}
|
|
|
|
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,
|
|
suite_kind: SuiteKind,
|
|
source_type: PySourceType,
|
|
) -> Option<DocstringStmt<'a>> {
|
|
// Notebooks don't have a concept of modules, therefore, don't recognise the first string as the module docstring.
|
|
if source_type.is_ipynb() && suite_kind == SuiteKind::TopLevel {
|
|
return None;
|
|
}
|
|
|
|
let Stmt::Expr(ast::StmtExpr { value, .. }) = stmt else {
|
|
return None;
|
|
};
|
|
|
|
match value.as_ref() {
|
|
Expr::StringLiteral(ast::ExprStringLiteral { value, .. })
|
|
if !value.is_implicit_concatenated() =>
|
|
{
|
|
Some(DocstringStmt {
|
|
docstring: stmt,
|
|
suite_kind,
|
|
})
|
|
}
|
|
_ => None,
|
|
}
|
|
}
|
|
|
|
pub(crate) fn is_docstring_statement(stmt: &StmtExpr, source_type: PySourceType) -> bool {
|
|
if source_type.is_ipynb() {
|
|
return false;
|
|
}
|
|
|
|
if let Expr::StringLiteral(ast::ExprStringLiteral { value, .. }) = stmt.value.as_ref() {
|
|
!value.is_implicit_concatenated()
|
|
} else {
|
|
false
|
|
}
|
|
}
|
|
}
|
|
|
|
impl Format<PyFormatContext<'_>> for DocstringStmt<'_> {
|
|
fn fmt(&self, f: &mut Formatter<PyFormatContext<'_>>) -> FormatResult<()> {
|
|
let comments = f.context().comments().clone();
|
|
let node_comments = comments.leading_dangling_trailing(self.docstring);
|
|
|
|
if FormatStmtExpr.is_suppressed(node_comments.trailing, f.context()) {
|
|
suppressed_node(self.docstring).fmt(f)
|
|
} else {
|
|
// SAFETY: Safe because `DocStringStmt` guarantees that it only ever wraps a `ExprStmt` containing a `ExprStringLiteral`.
|
|
let string_literal = self
|
|
.docstring
|
|
.as_expr_stmt()
|
|
.unwrap()
|
|
.value
|
|
.as_string_literal_expr()
|
|
.unwrap();
|
|
|
|
// We format the expression, but the statement carries the comments
|
|
write!(
|
|
f,
|
|
[
|
|
leading_comments(node_comments.leading),
|
|
f.options()
|
|
.source_map_generation()
|
|
.is_enabled()
|
|
.then_some(source_position(self.docstring.start())),
|
|
string_literal
|
|
.format()
|
|
.with_options(ExprStringLiteralKind::Docstring),
|
|
f.options()
|
|
.source_map_generation()
|
|
.is_enabled()
|
|
.then_some(source_position(self.docstring.end())),
|
|
]
|
|
)?;
|
|
|
|
if self.suite_kind == SuiteKind::Class {
|
|
// Comments after class docstrings need a newline between the docstring and the
|
|
// comment (https://github.com/astral-sh/ruff/issues/7948).
|
|
// ```python
|
|
// class ModuleBrowser:
|
|
// """Browse module classes and functions in IDLE."""
|
|
// # ^ Insert a newline above here
|
|
//
|
|
// def __init__(self, master, path, *, _htest=False, _utest=False):
|
|
// pass
|
|
// ```
|
|
if let Some(own_line) = node_comments
|
|
.trailing
|
|
.iter()
|
|
.find(|comment| comment.line_position().is_own_line())
|
|
{
|
|
if lines_before(own_line.start(), f.context().source()) < 2 {
|
|
empty_line().fmt(f)?;
|
|
}
|
|
}
|
|
}
|
|
|
|
trailing_comments(node_comments.trailing).fmt(f)
|
|
}
|
|
}
|
|
}
|
|
|
|
/// A Child of a suite.
|
|
#[derive(Copy, Clone, Debug)]
|
|
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.docstring,
|
|
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;
|
|
use ruff_python_parser::parse_module;
|
|
use ruff_python_trivia::CommentRanges;
|
|
|
|
use crate::comments::Comments;
|
|
use crate::prelude::*;
|
|
use crate::statement::suite::SuiteKind;
|
|
use crate::PyFormatOptions;
|
|
|
|
fn format_suite(level: SuiteKind) -> String {
|
|
let source = r"
|
|
a = 10
|
|
|
|
|
|
|
|
three_leading_newlines = 80
|
|
|
|
|
|
two_leading_newlines = 20
|
|
|
|
one_leading_newline = 10
|
|
no_leading_newline = 30
|
|
class InTheMiddle:
|
|
pass
|
|
trailing_statement = 1
|
|
def func():
|
|
pass
|
|
def trailing_func():
|
|
pass
|
|
";
|
|
|
|
let parsed = parse_module(source).unwrap();
|
|
let comment_ranges = CommentRanges::from(parsed.tokens());
|
|
|
|
let context = PyFormatContext::new(
|
|
PyFormatOptions::default(),
|
|
source,
|
|
Comments::from_ranges(&comment_ranges),
|
|
parsed.tokens(),
|
|
);
|
|
|
|
let test_formatter =
|
|
format_with(|f: &mut PyFormatter| parsed.suite().format().with_options(level).fmt(f));
|
|
|
|
let formatted = format!(context, [test_formatter]).unwrap();
|
|
let printed = formatted.print().unwrap();
|
|
|
|
printed.as_code().to_string()
|
|
}
|
|
|
|
#[test]
|
|
fn top_level() {
|
|
let formatted = format_suite(SuiteKind::TopLevel);
|
|
|
|
assert_eq!(
|
|
formatted,
|
|
r"a = 10
|
|
|
|
|
|
three_leading_newlines = 80
|
|
|
|
|
|
two_leading_newlines = 20
|
|
|
|
one_leading_newline = 10
|
|
no_leading_newline = 30
|
|
|
|
|
|
class InTheMiddle:
|
|
pass
|
|
|
|
|
|
trailing_statement = 1
|
|
|
|
|
|
def func():
|
|
pass
|
|
|
|
|
|
def trailing_func():
|
|
pass
|
|
"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn nested_level() {
|
|
let formatted = format_suite(SuiteKind::Other {
|
|
last_suite_in_statement: true,
|
|
});
|
|
|
|
assert_eq!(
|
|
formatted,
|
|
r"a = 10
|
|
|
|
three_leading_newlines = 80
|
|
|
|
two_leading_newlines = 20
|
|
|
|
one_leading_newline = 10
|
|
no_leading_newline = 30
|
|
|
|
class InTheMiddle:
|
|
pass
|
|
|
|
trailing_statement = 1
|
|
|
|
def func():
|
|
pass
|
|
|
|
def trailing_func():
|
|
pass
|
|
"
|
|
);
|
|
}
|
|
}
|