Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Format function and class definitions into a single line if its body is an ellipsis #6592

Merged
merged 15 commits into from
Aug 21, 2023
Original file line number Diff line number Diff line change
@@ -0,0 +1,50 @@
"""Compound statements with no body should be written on one line."""

while True:
...

if True:
# comment
...

with True:
... # comment

for i in []:
...

for i in []:
# comment
...

for i in []:
... # comment

if True:
...

if True:
# comment
...

if True:
... # comment

with True:
...

with True:
# comment
...

with True:
... # comment

match x:
case 1:
...
case 2:
# comment
...
case 3:
... # comment
Original file line number Diff line number Diff line change
Expand Up @@ -76,3 +76,16 @@ class Test2(A):
def b(): ...
# comment
def c(): ...

class EllipsisWithComment:
... # comment

tjkuson marked this conversation as resolved.
Show resolved Hide resolved
def function_with_comment():
... # comment

class EllispsisWithMultipleTrailing: # trailing class comment
... # trailing ellipsis comment

class EllipsisWithLeadingComment:
# leading
...
56 changes: 54 additions & 2 deletions crates/ruff_python_formatter/src/statement/clause.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,12 +2,13 @@ use crate::comments::{
leading_alternate_branch_comments, trailing_comments, SourceComment, SuppressionKind,
};
use crate::prelude::*;
use crate::statement::suite::{contains_only_an_ellipsis, SuiteKind};
use crate::verbatim::write_suppressed_clause_header;
use ruff_formatter::{Argument, Arguments, FormatError};
use ruff_formatter::{write, Argument, Arguments, FormatError};
use ruff_python_ast::node::AnyNodeRef;
use ruff_python_ast::{
ElifElseClause, ExceptHandlerExceptHandler, MatchCase, Ranged, StmtClassDef, StmtFor,
StmtFunctionDef, StmtIf, StmtMatch, StmtTry, StmtWhile, StmtWith,
StmtFunctionDef, StmtIf, StmtMatch, StmtTry, StmtWhile, StmtWith, Suite,
};
use ruff_python_trivia::{SimpleToken, SimpleTokenKind, SimpleTokenizer};
use ruff_text_size::{TextRange, TextSize};
Expand Down Expand Up @@ -352,6 +353,57 @@ impl<'ast> Format<PyFormatContext<'ast>> for FormatClauseHeader<'_, 'ast> {
}
}

pub(crate) struct FormatClauseBody<'a> {
body: &'a Suite,
kind: SuiteKind,
trailing_comments: &'a [SourceComment],
}

impl<'a> FormatClauseBody<'a> {
#[must_use]
pub(crate) fn with_kind(mut self, kind: SuiteKind) -> Self {
self.kind = kind;
self
}
}

pub(crate) fn clause_body<'a>(
body: &'a Suite,
trailing_comments: &'a [SourceComment],
) -> FormatClauseBody<'a> {
FormatClauseBody {
body,
kind: SuiteKind::default(),
trailing_comments,
}
}

impl Format<PyFormatContext<'_>> for FormatClauseBody<'_> {
fn fmt(&self, f: &mut Formatter<PyFormatContext<'_>>) -> FormatResult<()> {
if f.options().source_type().is_stub()
&& contains_only_an_ellipsis(self.body, f.context().comments())
&& self.trailing_comments.is_empty()
{
write!(
f,
[
space(),
self.body.format().with_options(self.kind),
hard_line_break()
]
)
} else {
write!(
f,
[
trailing_comments(self.trailing_comments),
block_indent(&self.body.format().with_options(self.kind))
]
)
}
}
}

/// Finds the range of `keyword` starting the search at `start_position`. Expects only comments and `(` between
/// the `start_position` and the `keyword` token.
fn find_keyword(
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@ use crate::comments::{leading_comments, trailing_comments, SourceComment};
use crate::prelude::*;
use crate::statement::suite::SuiteKind;

use crate::statement::clause::{clause_header, ClauseHeader};
use crate::statement::clause::{clause_body, clause_header, ClauseHeader};
use crate::FormatNodeRule;

#[derive(Default)]
Expand Down Expand Up @@ -107,7 +107,7 @@ impl FormatNodeRule<StmtClassDef> for FormatStmtClassDef {
Ok(())
}),
),
block_indent(&body.format().with_options(SuiteKind::Class))
clause_body(body, trailing_definition_comments).with_kind(SuiteKind::Class),
]
)
}
Expand Down
4 changes: 2 additions & 2 deletions crates/ruff_python_formatter/src/statement/stmt_for.rs
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@ use crate::expression::expr_tuple::TupleParentheses;
use crate::expression::maybe_parenthesize_expression;
use crate::expression::parentheses::Parenthesize;
use crate::prelude::*;
use crate::statement::clause::{clause_header, ClauseHeader, ElseClause};
use crate::statement::clause::{clause_body, clause_header, ClauseHeader, ElseClause};
use crate::FormatNodeRule;

#[derive(Debug)]
Expand Down Expand Up @@ -64,7 +64,7 @@ impl FormatNodeRule<StmtFor> for FormatStmtFor {
maybe_parenthesize_expression(iter, item, Parenthesize::IfBreaks),
],
),
block_indent(&body.format())
clause_body(body, trailing_condition_comments),
tjkuson marked this conversation as resolved.
Show resolved Hide resolved
]
)?;

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@ use crate::comments::SourceComment;
use crate::expression::maybe_parenthesize_expression;
use crate::expression::parentheses::{Parentheses, Parenthesize};
use crate::prelude::*;
use crate::statement::clause::{clause_header, ClauseHeader};
use crate::statement::clause::{clause_body, clause_header, ClauseHeader};
use crate::statement::stmt_class_def::FormatDecorators;
use crate::statement::suite::SuiteKind;
use crate::FormatNodeRule;
Expand Down Expand Up @@ -142,7 +142,7 @@ impl FormatNodeRule<StmtFunctionDef> for FormatStmtFunctionDef {
group(&format_inner).fmt(f)
}),
),
block_indent(&body.format().with_options(SuiteKind::Function))
clause_body(body, trailing_definition_comments).with_kind(SuiteKind::Function),
]
)
}
Expand Down
4 changes: 2 additions & 2 deletions crates/ruff_python_formatter/src/statement/stmt_if.rs
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@ use crate::comments::SourceComment;
use crate::expression::maybe_parenthesize_expression;
use crate::expression::parentheses::Parenthesize;
use crate::prelude::*;
use crate::statement::clause::{clause_header, ClauseHeader};
use crate::statement::clause::{clause_body, clause_header, ClauseHeader};
use crate::FormatNodeRule;

#[derive(Default)]
Expand Down Expand Up @@ -36,7 +36,7 @@ impl FormatNodeRule<StmtIf> for FormatStmtIf {
maybe_parenthesize_expression(test, item, Parenthesize::IfBreaks),
],
),
block_indent(&body.format())
clause_body(body, trailing_colon_comment),
tjkuson marked this conversation as resolved.
Show resolved Hide resolved
]
)?;

Expand Down
4 changes: 2 additions & 2 deletions crates/ruff_python_formatter/src/statement/stmt_while.rs
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@ use crate::comments::SourceComment;
use crate::expression::maybe_parenthesize_expression;
use crate::expression::parentheses::Parenthesize;
use crate::prelude::*;
use crate::statement::clause::{clause_header, ClauseHeader, ElseClause};
use crate::statement::clause::{clause_body, clause_header, ClauseHeader, ElseClause};
use crate::FormatNodeRule;

#[derive(Default)]
Expand Down Expand Up @@ -43,7 +43,7 @@ impl FormatNodeRule<StmtWhile> for FormatStmtWhile {
maybe_parenthesize_expression(test, item, Parenthesize::IfBreaks),
]
),
block_indent(&body.format())
clause_body(body, trailing_condition_comments),
tjkuson marked this conversation as resolved.
Show resolved Hide resolved
]
)?;

Expand Down
4 changes: 2 additions & 2 deletions crates/ruff_python_formatter/src/statement/stmt_with.rs
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@ use crate::expression::parentheses::{
};
use crate::other::commas;
use crate::prelude::*;
use crate::statement::clause::{clause_header, ClauseHeader};
use crate::statement::clause::{clause_body, clause_header, ClauseHeader};
use crate::{FormatNodeRule, PyFormatOptions};

#[derive(Default)]
Expand Down Expand Up @@ -100,7 +100,7 @@ impl FormatNodeRule<StmtWith> for FormatStmtWith {
Ok(())
})
),
block_indent(&item.body.format())
clause_body(&item.body, colon_comments)
]
)
}
Expand Down
41 changes: 27 additions & 14 deletions crates/ruff_python_formatter/src/statement/suite.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
use crate::comments::{leading_comments, trailing_comments};
use crate::comments::{leading_comments, trailing_comments, Comments};
use ruff_formatter::{write, FormatOwnedWithRule, FormatRefWithRule, FormatRuleWithOptions};
use ruff_python_ast::helpers::is_compound_statement;
use ruff_python_ast::node::AnyNodeRef;
Expand All @@ -17,7 +17,7 @@ use crate::verbatim::{
};

/// Level at which the [`Suite`] appears in the source code.
#[derive(Copy, Clone, Debug)]
#[derive(Copy, Clone, Debug, Default)]
pub enum SuiteKind {
/// Statements at the module level / top level
TopLevel,
Expand All @@ -29,6 +29,7 @@ pub enum SuiteKind {
Class,

/// Statements in any other body (e.g., `if` or `while`).
#[default]
Other,
}

Expand Down Expand Up @@ -168,9 +169,15 @@ impl FormatRule<Suite, PyFormatContext<'_>> for FormatSuite {
// ```
let class_sequences_with_ellipsis_only =
preceding.as_class_def_stmt().is_some_and(|class| {
contains_only_an_ellipsis(&class.body)
contains_only_an_ellipsis(
&class.body,
f.context().comments(),
)
}) && following.as_class_def_stmt().is_some_and(|class| {
contains_only_an_ellipsis(&class.body)
contains_only_an_ellipsis(
&class.body,
f.context().comments(),
)
});

// Two subsequent functions where the preceding has an ellipsis only body
Expand All @@ -180,7 +187,10 @@ impl FormatRule<Suite, PyFormatContext<'_>> for FormatSuite {
// ```
let function_with_ellipsis =
preceding.as_function_def_stmt().is_some_and(|function| {
contains_only_an_ellipsis(&function.body)
contains_only_an_ellipsis(
&function.body,
f.context().comments(),
)
}) && following.is_function_def_stmt();

// Don't add an empty line between two classes that have an `...` body only or after
Expand Down Expand Up @@ -325,16 +335,19 @@ impl FormatRule<Suite, PyFormatContext<'_>> for FormatSuite {
}
}

/// Returns `true` if a function or class body contains only an ellipsis.
fn contains_only_an_ellipsis(body: &[Stmt]) -> bool {
/// 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, .. })] => matches!(
value.as_ref(),
Expr::Constant(ast::ExprConstant {
value: Constant::Ellipsis,
..
})
),
[Stmt::Expr(ast::StmtExpr { value, .. })] => {
let [node] = body else { return false; };
matches!(
value.as_ref(),
Expr::Constant(ast::ExprConstant {
value: Constant::Ellipsis,
..
})
) && !comments.has_leading(node)
}
_ => false,
}
}
Expand Down