use tracing::Level;
use ruff_formatter::printer::SourceMapGeneration;
use ruff_formatter::{
FormatContext, FormatError, FormatOptions, IndentStyle, PrintedRange, SourceCode, format,
};
use ruff_python_ast::visitor::source_order::{SourceOrderVisitor, TraversalSignal, walk_body};
use ruff_python_ast::{AnyNodeRef, Stmt, StmtMatch, StmtTry};
use ruff_python_parser::{ParseOptions, parse};
use ruff_python_trivia::{
BackwardsTokenizer, SimpleToken, SimpleTokenKind, TriviaRanges, indentation_at_offset,
};
use ruff_text_size::{Ranged, TextLen, TextRange, TextSize};
use crate::comments::Comments;
use crate::context::{IndentLevel, NodeLevel};
use crate::prelude::*;
use crate::statement::suite::{DocstringStmt, skip_range};
use crate::verbatim::{ends_suppression, starts_suppression};
use crate::{FormatModuleError, PyFormatOptions, format_module_source};
#[tracing::instrument(name = "format_range", level = Level::TRACE, skip_all)]
pub fn format_range(
source: &str,
range: TextRange,
options: PyFormatOptions,
) -> Result<PrintedRange, FormatModuleError> {
if source.text_len() < range.end() {
return Err(FormatModuleError::FormatError(FormatError::RangeError {
input: range,
tree: TextRange::up_to(source.text_len()),
}));
}
if range.is_empty() {
return Ok(PrintedRange::empty());
}
if range == TextRange::up_to(source.text_len()) {
let formatted = format_module_source(source, options)?;
return Ok(PrintedRange::new(formatted.into_code(), range));
}
assert_valid_char_boundaries(range, source);
let parsed = parse(source, ParseOptions::from(options.source_type()))?;
let source_code = SourceCode::new(source);
let trivia = TriviaRanges::from(parsed.tokens());
let comments = Comments::from_ast(parsed.syntax(), source_code, &trivia);
let mut context = PyFormatContext::new(
options.with_source_map_generation(SourceMapGeneration::Enabled),
source,
comments,
&trivia,
parsed.tokens(),
);
let (enclosing_node, base_indent) =
match find_enclosing_node(range, AnyNodeRef::from(parsed.syntax()), &context) {
EnclosingNode::Node { node, indent_level } => (node, indent_level),
EnclosingNode::Suppressed => {
return Ok(PrintedRange::empty());
}
};
let narrowed_range = narrow_range(range, enclosing_node, &context);
assert_valid_char_boundaries(narrowed_range, source);
if !enclosing_node.is_mod_module() {
context.set_node_level(NodeLevel::CompoundStatement);
context.set_indent_level(
IndentLevel::new(base_indent.saturating_add(1)),
);
}
let formatted = format!(
context,
[FormatEnclosingNode {
root: enclosing_node
}]
)?;
let printed = formatted.print_with_indent(base_indent)?;
Ok(printed.slice_range(narrowed_range, source))
}
fn find_enclosing_node<'ast>(
range: TextRange,
root: AnyNodeRef<'ast>,
context: &PyFormatContext<'ast>,
) -> EnclosingNode<'ast> {
let mut visitor = FindEnclosingNode::new(range, context);
if visitor.enter_node(root).is_traverse() {
root.visit_source_order(&mut visitor);
}
visitor.leave_node(root);
visitor.closest
}
struct FindEnclosingNode<'a, 'ast> {
range: TextRange,
context: &'a PyFormatContext<'ast>,
closest: EnclosingNode<'ast>,
suppressed: Suppressed,
}
impl<'a, 'ast> FindEnclosingNode<'a, 'ast> {
fn new(range: TextRange, context: &'a PyFormatContext<'ast>) -> Self {
Self {
range,
context,
suppressed: Suppressed::No,
closest: EnclosingNode::Suppressed,
}
}
}
impl<'ast> SourceOrderVisitor<'ast> for FindEnclosingNode<'_, 'ast> {
fn enter_node(&mut self, node: AnyNodeRef<'ast>) -> TraversalSignal {
if !(is_logical_line(node) || node.is_mod_module()) {
return TraversalSignal::Skip;
}
if node.is_statement() {
let leading_comments = self.context.comments().leading(node);
self.suppressed = Suppressed::from(match self.suppressed {
Suppressed::No => starts_suppression(leading_comments, self.context.source()),
Suppressed::Yes => !ends_suppression(leading_comments, self.context.source()),
});
}
if !node.range().contains_range(self.range) {
return TraversalSignal::Skip;
}
if self.suppressed.is_yes() && node.is_statement() {
self.closest = EnclosingNode::Suppressed;
return TraversalSignal::Skip;
}
let is_maybe_docstring = node
.as_stmt_expr()
.is_some_and(|stmt| DocstringStmt::is_docstring_statement(stmt, self.context));
if is_maybe_docstring {
return TraversalSignal::Skip;
}
let Some(indent_level) =
indent_level(node.start(), self.context.source(), self.context.options())
else {
return TraversalSignal::Skip;
};
self.closest = EnclosingNode::Node { node, indent_level };
TraversalSignal::Traverse
}
fn leave_node(&mut self, node: AnyNodeRef<'ast>) {
if node.is_statement() {
let trailing_comments = self.context.comments().trailing(node);
self.suppressed = Suppressed::from(match self.suppressed {
Suppressed::No => starts_suppression(trailing_comments, self.context.source()),
Suppressed::Yes => !ends_suppression(trailing_comments, self.context.source()),
});
}
}
fn visit_body(&mut self, body: &'ast [Stmt]) {
debug_assert!(self.suppressed.is_no());
let mut iter = body.iter();
while let Some(stmt) = iter.next() {
if let Some(verbatim_range) = skip_range(stmt, iter.as_slice(), self.context)
&& verbatim_range.intersect(self.range).is_some()
{
break;
}
self.visit_stmt(stmt);
}
self.suppressed = Suppressed::No;
}
}
#[derive(Debug, Copy, Clone)]
enum EnclosingNode<'a> {
Suppressed,
Node {
node: AnyNodeRef<'a>,
indent_level: u16,
},
}
fn narrow_range(
range: TextRange,
enclosing_node: AnyNodeRef,
context: &PyFormatContext,
) -> TextRange {
let enclosing_indent = indentation_at_offset(enclosing_node.start(), context.source())
.expect("Expected enclosing to never be a same line body statement.");
let mut visitor = NarrowRange {
context,
range,
narrowed_start: enclosing_node.start(),
narrowed_end: enclosing_node.end(),
enclosing_indent,
level: usize::from(!enclosing_node.is_mod_module()),
};
if visitor.enter_node(enclosing_node).is_traverse() {
enclosing_node.visit_source_order(&mut visitor);
}
visitor.leave_node(enclosing_node);
TextRange::new(visitor.narrowed_start, visitor.narrowed_end)
}
struct NarrowRange<'a> {
context: &'a PyFormatContext<'a>,
range: TextRange,
narrowed_start: TextSize,
narrowed_end: TextSize,
enclosing_indent: &'a str,
level: usize,
}
impl SourceOrderVisitor<'_> for NarrowRange<'_> {
fn enter_node(&mut self, node: AnyNodeRef<'_>) -> TraversalSignal {
if !(is_logical_line(node) || node.is_mod_module()) {
return TraversalSignal::Skip;
}
let leading_comments = self.context.comments().leading(node);
self.narrow(leading_comments);
self.narrow([node]);
if node.end() < self.range.start()
|| (self.narrowed_start > node.start() && self.narrowed_end <= node.end())
{
return TraversalSignal::Skip;
}
match node {
AnyNodeRef::StmtMatch(StmtMatch {
subject: _,
cases,
range: _,
node_index: _,
}) => {
if let Some(saved_state) = self.enter_level(cases.first().map(AnyNodeRef::from)) {
for match_case in cases {
self.visit_match_case(match_case);
}
self.leave_level(saved_state);
}
TraversalSignal::Skip
}
AnyNodeRef::StmtTry(StmtTry {
body,
handlers,
orelse,
finalbody,
is_star: _,
range: _,
node_index: _,
}) => {
self.visit_body(body);
if let Some(except_handler_saved) =
self.enter_level(handlers.first().map(AnyNodeRef::from))
{
for except_handler in handlers {
self.visit_except_handler(except_handler);
}
self.leave_level(except_handler_saved);
}
self.visit_body(orelse);
self.visit_body(finalbody);
TraversalSignal::Skip
}
_ => TraversalSignal::Traverse,
}
}
fn leave_node(&mut self, node: AnyNodeRef<'_>) {
if !(is_logical_line(node) || node.is_mod_module()) {
return;
}
self.narrow(
self.context
.comments()
.trailing(node)
.iter()
.filter(|comment| comment.line_position().is_own_line()),
);
}
fn visit_body(&mut self, body: &'_ [Stmt]) {
if let Some(saved_state) = self.enter_level(body.first().map(AnyNodeRef::from)) {
walk_body(self, body);
self.leave_level(saved_state);
}
}
}
impl NarrowRange<'_> {
fn narrow<I, T>(&mut self, items: I)
where
I: IntoIterator<Item = T>,
T: Ranged,
{
for ranged in items {
self.narrow_offset(ranged.start());
self.narrow_offset(ranged.end());
}
}
fn narrow_offset(&mut self, offset: TextSize) {
self.narrow_start(offset);
self.narrow_end(offset);
}
fn narrow_start(&mut self, offset: TextSize) {
if offset <= self.range.start() {
self.narrowed_start = self.narrowed_start.max(offset);
}
}
fn narrow_end(&mut self, offset: TextSize) {
if offset >= self.range.end() {
self.narrowed_end = self.narrowed_end.min(offset);
}
}
fn enter_level(&mut self, first_child: Option<AnyNodeRef>) -> Option<SavedLevel> {
if let Some(first_child) = first_child {
if let Some(SimpleToken {
kind: SimpleTokenKind::Colon,
range: colon_range,
}) = BackwardsTokenizer::up_to(
first_child.start(),
self.context.source(),
self.context.trivia().comments(),
)
.skip_trivia()
.next()
{
self.narrow_offset(colon_range.end());
}
if let Some(indentation) =
indentation_at_offset(first_child.start(), self.context.source())
{
let relative_indent = indentation.strip_prefix(self.enclosing_indent).unwrap();
let expected_indents = self.level;
let has_expected_indentation = match self.context.options().indent_style() {
IndentStyle::Tab => {
relative_indent.len() == expected_indents
&& relative_indent.chars().all(|c| c == '\t')
}
IndentStyle::Space => {
relative_indent.len()
== expected_indents
* self.context.options().indent_width().value() as usize
&& relative_indent.chars().all(|c| c == ' ')
}
};
if !has_expected_indentation {
return None;
}
} else {
return None;
}
}
let saved_level = self.level;
self.level += 1;
Some(SavedLevel { level: saved_level })
}
#[expect(clippy::needless_pass_by_value)]
fn leave_level(&mut self, saved_state: SavedLevel) {
self.level = saved_state.level;
}
}
pub(crate) const fn is_logical_line(node: AnyNodeRef) -> bool {
node.is_statement()
|| node.is_decorator()
|| node.is_except_handler()
|| node.is_elif_else_clause()
|| node.is_match_case()
}
#[derive(Debug)]
struct SavedLevel {
level: usize,
}
#[derive(Copy, Clone, Default, Debug)]
enum Suppressed {
#[default]
No,
Yes,
}
impl Suppressed {
const fn is_no(self) -> bool {
matches!(self, Suppressed::No)
}
const fn is_yes(self) -> bool {
matches!(self, Suppressed::Yes)
}
}
impl From<bool> for Suppressed {
fn from(value: bool) -> Self {
if value {
Suppressed::Yes
} else {
Suppressed::No
}
}
}
fn assert_valid_char_boundaries(range: TextRange, source: &str) {
assert!(source.is_char_boundary(usize::from(range.start())));
assert!(source.is_char_boundary(usize::from(range.end())));
}
struct FormatEnclosingNode<'a> {
root: AnyNodeRef<'a>,
}
impl Format<PyFormatContext<'_>> for FormatEnclosingNode<'_> {
fn fmt(&self, f: &mut Formatter<PyFormatContext<'_>>) -> FormatResult<()> {
match self.root {
AnyNodeRef::ModModule(node) => node.format().fmt(f),
AnyNodeRef::ModExpression(node) => node.format().fmt(f),
AnyNodeRef::StmtFunctionDef(node) => node.format().fmt(f),
AnyNodeRef::StmtClassDef(node) => node.format().fmt(f),
AnyNodeRef::StmtReturn(node) => node.format().fmt(f),
AnyNodeRef::StmtDelete(node) => node.format().fmt(f),
AnyNodeRef::StmtTypeAlias(node) => node.format().fmt(f),
AnyNodeRef::StmtAssign(node) => node.format().fmt(f),
AnyNodeRef::StmtAugAssign(node) => node.format().fmt(f),
AnyNodeRef::StmtAnnAssign(node) => node.format().fmt(f),
AnyNodeRef::StmtFor(node) => node.format().fmt(f),
AnyNodeRef::StmtWhile(node) => node.format().fmt(f),
AnyNodeRef::StmtIf(node) => node.format().fmt(f),
AnyNodeRef::StmtWith(node) => node.format().fmt(f),
AnyNodeRef::StmtMatch(node) => node.format().fmt(f),
AnyNodeRef::StmtRaise(node) => node.format().fmt(f),
AnyNodeRef::StmtTry(node) => node.format().fmt(f),
AnyNodeRef::StmtAssert(node) => node.format().fmt(f),
AnyNodeRef::StmtImport(node) => node.format().fmt(f),
AnyNodeRef::StmtImportFrom(node) => node.format().fmt(f),
AnyNodeRef::StmtGlobal(node) => node.format().fmt(f),
AnyNodeRef::StmtNonlocal(node) => node.format().fmt(f),
AnyNodeRef::StmtExpr(node) => node.format().fmt(f),
AnyNodeRef::StmtPass(node) => node.format().fmt(f),
AnyNodeRef::StmtBreak(node) => node.format().fmt(f),
AnyNodeRef::StmtContinue(node) => node.format().fmt(f),
AnyNodeRef::StmtIpyEscapeCommand(node) => node.format().fmt(f),
AnyNodeRef::ExceptHandlerExceptHandler(node) => node.format().fmt(f),
AnyNodeRef::MatchCase(node) => node.format().fmt(f),
AnyNodeRef::Decorator(node) => node.format().fmt(f),
AnyNodeRef::ElifElseClause(node) => node.format().fmt(f),
AnyNodeRef::ExprBoolOp(_)
| AnyNodeRef::ExprNamed(_)
| AnyNodeRef::ExprBinOp(_)
| AnyNodeRef::ExprUnaryOp(_)
| AnyNodeRef::ExprLambda(_)
| AnyNodeRef::ExprIf(_)
| AnyNodeRef::ExprDict(_)
| AnyNodeRef::ExprSet(_)
| AnyNodeRef::ExprListComp(_)
| AnyNodeRef::ExprSetComp(_)
| AnyNodeRef::ExprDictComp(_)
| AnyNodeRef::ExprGenerator(_)
| AnyNodeRef::ExprAwait(_)
| AnyNodeRef::ExprYield(_)
| AnyNodeRef::ExprYieldFrom(_)
| AnyNodeRef::ExprCompare(_)
| AnyNodeRef::ExprCall(_)
| AnyNodeRef::InterpolatedElement(_)
| AnyNodeRef::InterpolatedStringLiteralElement(_)
| AnyNodeRef::InterpolatedStringFormatSpec(_)
| AnyNodeRef::ExprFString(_)
| AnyNodeRef::ExprTString(_)
| AnyNodeRef::ExprStringLiteral(_)
| AnyNodeRef::ExprBytesLiteral(_)
| AnyNodeRef::ExprNumberLiteral(_)
| AnyNodeRef::ExprBooleanLiteral(_)
| AnyNodeRef::ExprNoneLiteral(_)
| AnyNodeRef::ExprEllipsisLiteral(_)
| AnyNodeRef::ExprAttribute(_)
| AnyNodeRef::ExprSubscript(_)
| AnyNodeRef::ExprStarred(_)
| AnyNodeRef::ExprName(_)
| AnyNodeRef::ExprList(_)
| AnyNodeRef::ExprTuple(_)
| AnyNodeRef::ExprSlice(_)
| AnyNodeRef::ExprIpyEscapeCommand(_)
| AnyNodeRef::FString(_)
| AnyNodeRef::StringLiteral(_)
| AnyNodeRef::TString(_)
| AnyNodeRef::PatternMatchValue(_)
| AnyNodeRef::PatternMatchSingleton(_)
| AnyNodeRef::PatternMatchSequence(_)
| AnyNodeRef::PatternMatchMapping(_)
| AnyNodeRef::PatternMatchClass(_)
| AnyNodeRef::PatternMatchStar(_)
| AnyNodeRef::PatternMatchAs(_)
| AnyNodeRef::PatternMatchOr(_)
| AnyNodeRef::PatternArguments(_)
| AnyNodeRef::PatternKeyword(_)
| AnyNodeRef::Comprehension(_)
| AnyNodeRef::Arguments(_)
| AnyNodeRef::Parameters(_)
| AnyNodeRef::Parameter(_)
| AnyNodeRef::ParameterWithDefault(_)
| AnyNodeRef::Keyword(_)
| AnyNodeRef::Alias(_)
| AnyNodeRef::WithItem(_)
| AnyNodeRef::TypeParams(_)
| AnyNodeRef::TypeParamTypeVar(_)
| AnyNodeRef::TypeParamTypeVarTuple(_)
| AnyNodeRef::TypeParamParamSpec(_)
| AnyNodeRef::Identifier(_)
| AnyNodeRef::BytesLiteral(_) => {
panic!("Range formatting only supports formatting logical lines")
}
}
}
}
fn indent_level(offset: TextSize, source: &str, options: &PyFormatOptions) -> Option<u16> {
let indentation = indentation_at_offset(offset, source)?;
let level = match options.indent_style() {
IndentStyle::Tab => {
if indentation.chars().all(|c| c == '\t') {
Some(indentation.len())
} else {
None
}
}
IndentStyle::Space => {
let indent_width = options.indent_width().value() as usize;
if indentation.chars().all(|c| c == ' ') && indentation.len() % indent_width == 0 {
Some(indentation.len() / indent_width)
} else {
None
}
}
};
level.map(|level| u16::try_from(level).unwrap_or(u16::MAX))
}