Preserve trailing semicolon for Notebooks (#8590)

## Summary

This PR updates the formatter to preserve trailing semicolon for Jupyter
Notebooks.

The motivation behind the change is that semicolons in notebooks are
typically used to hide the output, for example when plotting. This is
highlighted in the linked issue.

The conditions required as to when the trailing semicolon should be
preserved are:
1. It should be a top-level statement which is last in the module.
2. For statement, it can be either assignment, annotated assignment, or
augmented assignment. Here, the target should only be a single
identifier i.e., multiple assignments or tuple unpacking isn't
considered.
3. For expression, it can be any.

## Test Plan

Add a new integration test in `ruff_cli`. The test notebook basically
acts as a document as to which trailing semicolons are to be preserved.

fixes: #8254

crates/ruff_python_formatter/src/comments/format.rs

@@ -323,7 +323,7 @@ pub(crate) struct FormatEmptyLines {
 impl Format<PyFormatContext<'_>> for FormatEmptyLines {
     fn fmt(&self, f: &mut Formatter<PyFormatContext>) -> FormatResult<()> {
         match f.context().node_level() {
-            NodeLevel::TopLevel => match self.lines {
+            NodeLevel::TopLevel(_) => match self.lines {
                 0 | 1 => write!(f, [hard_line_break()]),
                 2 => write!(f, [empty_line()]),
                 _ => match f.options().source_type() {
@@ -519,9 +519,9 @@ pub(crate) fn empty_lines_before_trailing_comments<'a>(
 ) -> FormatEmptyLinesBeforeTrailingComments<'a> {
     // Black has different rules for stub vs. non-stub and top level vs. indented
     let empty_lines = match (f.options().source_type(), f.context().node_level()) {
-        (PySourceType::Stub, NodeLevel::TopLevel) => 1,
+        (PySourceType::Stub, NodeLevel::TopLevel(_)) => 1,
         (PySourceType::Stub, _) => 0,
-        (_, NodeLevel::TopLevel) => 2,
+        (_, NodeLevel::TopLevel(_)) => 2,
         (_, _) => 1,
     };
 
@@ -573,9 +573,9 @@ pub(crate) fn empty_lines_after_leading_comments<'a>(
 ) -> FormatEmptyLinesAfterLeadingComments<'a> {
     // Black has different rules for stub vs. non-stub and top level vs. indented
     let empty_lines = match (f.options().source_type(), f.context().node_level()) {
-        (PySourceType::Stub, NodeLevel::TopLevel) => 1,
+        (PySourceType::Stub, NodeLevel::TopLevel(_)) => 1,
         (PySourceType::Stub, _) => 0,
-        (_, NodeLevel::TopLevel) => 2,
+        (_, NodeLevel::TopLevel(_)) => 2,
         (_, _) => 1,
     };
 

crates/ruff_python_formatter/src/context.rs

@@ -19,7 +19,7 @@ impl<'a> PyFormatContext<'a> {
             options,
             contents,
             comments,
-            node_level: NodeLevel::TopLevel,
+            node_level: NodeLevel::TopLevel(TopLevelStatementPosition::Other),
         }
     }
 
@@ -68,12 +68,21 @@ impl Debug for PyFormatContext<'_> {
     }
 }
 
-/// What's the enclosing level of the outer node.
+/// The position of a top-level statement in the module.
 #[derive(Copy, Clone, Debug, Eq, PartialEq, Default)]
+pub(crate) enum TopLevelStatementPosition {
+    /// This is the last top-level statement in the module.
+    Last,
+    /// Any other top-level statement.
+    #[default]
+    Other,
+}
+
+/// What's the enclosing level of the outer node.
+#[derive(Copy, Clone, Debug, Eq, PartialEq)]
 pub(crate) enum NodeLevel {
     /// Formatting statements on the module level.
-    #[default]
-    TopLevel,
+    TopLevel(TopLevelStatementPosition),
 
     /// Formatting the body statements of a [compound statement](https://docs.python.org/3/reference/compound_stmts.html#compound-statements)
     /// (`if`, `while`, `match`, etc.).
@@ -86,6 +95,12 @@ pub(crate) enum NodeLevel {
     ParenthesizedExpression,
 }
 
+impl Default for NodeLevel {
+    fn default() -> Self {
+        Self::TopLevel(TopLevelStatementPosition::Other)
+    }
+}
+
 impl NodeLevel {
     /// Returns `true` if the expression is in a parenthesized context.
     pub(crate) const fn is_parenthesized(self) -> bool {
@@ -94,6 +109,11 @@ impl NodeLevel {
             NodeLevel::Expression(Some(_)) | NodeLevel::ParenthesizedExpression
         )
     }
+
+    /// Returns `true` if this is the last top-level statement in the module.
+    pub(crate) const fn is_last_top_level_statement(self) -> bool {
+        matches!(self, NodeLevel::TopLevel(TopLevelStatementPosition::Last))
+    }
 }
 
 /// Change the [`NodeLevel`] of the formatter for the lifetime of this struct

crates/ruff_python_formatter/src/expression/mod.rs

@@ -135,7 +135,9 @@ impl FormatRule<Expr, PyFormatContext<'_>> for FormatExpr {
             }
         } else {
             let level = match f.context().node_level() {
-                NodeLevel::TopLevel | NodeLevel::CompoundStatement => NodeLevel::Expression(None),
+                NodeLevel::TopLevel(_) | NodeLevel::CompoundStatement => {
+                    NodeLevel::Expression(None)
+                }
                 saved_level @ (NodeLevel::Expression(_) | NodeLevel::ParenthesizedExpression) => {
                     saved_level
                 }

crates/ruff_python_formatter/src/expression/parentheses.rs

@@ -252,7 +252,7 @@ pub(crate) enum InParenthesesOnlyLineBreak {
 impl<'ast> Format<PyFormatContext<'ast>> for InParenthesesOnlyLineBreak {
     fn fmt(&self, f: &mut Formatter<PyFormatContext<'ast>>) -> FormatResult<()> {
         match f.context().node_level() {
-            NodeLevel::TopLevel | NodeLevel::CompoundStatement | NodeLevel::Expression(None) => {
+            NodeLevel::TopLevel(_) | NodeLevel::CompoundStatement | NodeLevel::Expression(None) => {
                 match self {
                     InParenthesesOnlyLineBreak::SoftLineBreak => Ok(()),
                     InParenthesesOnlyLineBreak::SoftLineBreakOrSpace => space().fmt(f),
@@ -319,7 +319,7 @@ pub(super) fn write_in_parentheses_only_group_start_tag(f: &mut PyFormatter) {
             // Unconditionally group the content if it is not enclosed by an optional parentheses group.
             f.write_element(FormatElement::Tag(tag::Tag::StartGroup(tag::Group::new())));
         }
-        NodeLevel::Expression(None) | NodeLevel::TopLevel | NodeLevel::CompoundStatement => {
+        NodeLevel::Expression(None) | NodeLevel::TopLevel(_) | NodeLevel::CompoundStatement => {
             // No group
         }
     }
@@ -334,7 +334,7 @@ pub(super) fn write_in_parentheses_only_group_end_tag(f: &mut PyFormatter) {
             // Unconditionally group the content if it is not enclosed by an optional parentheses group.
             f.write_element(FormatElement::Tag(tag::Tag::EndGroup));
         }
-        NodeLevel::Expression(None) | NodeLevel::TopLevel | NodeLevel::CompoundStatement => {
+        NodeLevel::Expression(None) | NodeLevel::TopLevel(_) | NodeLevel::CompoundStatement => {
             // No group
         }
     }

crates/ruff_python_formatter/src/statement/stmt_ann_assign.rs

@@ -2,9 +2,11 @@ use ruff_formatter::write;
 use ruff_python_ast::StmtAnnAssign;
 
 use crate::comments::{SourceComment, SuppressionKind};
+
 use crate::expression::maybe_parenthesize_expression;
 use crate::expression::parentheses::Parenthesize;
 use crate::prelude::*;
+use crate::statement::trailing_semicolon;
 
 #[derive(Default)]
 pub struct FormatStmtAnnAssign;
@@ -36,6 +38,14 @@ impl FormatNodeRule<StmtAnnAssign> for FormatStmtAnnAssign {
             )?;
         }
 
+        if f.options().source_type().is_ipynb()
+            && f.context().node_level().is_last_top_level_statement()
+            && target.is_name_expr()
+            && trailing_semicolon(item.into(), f.context().source()).is_some()
+        {
+            token(";").fmt(f)?;
+        }
+
         Ok(())
     }
 

crates/ruff_python_formatter/src/statement/stmt_assign.rs

@@ -6,6 +6,7 @@ use crate::context::{NodeLevel, WithNodeLevel};
 use crate::expression::parentheses::{Parentheses, Parenthesize};
 use crate::expression::{has_own_parentheses, maybe_parenthesize_expression};
 use crate::prelude::*;
+use crate::statement::trailing_semicolon;
 
 #[derive(Default)]
 pub struct FormatStmtAssign;
@@ -40,7 +41,18 @@ impl FormatNodeRule<StmtAssign> for FormatStmtAssign {
                 item,
                 Parenthesize::IfBreaks
             )]
-        )
+        )?;
+
+        if f.options().source_type().is_ipynb()
+            && f.context().node_level().is_last_top_level_statement()
+            && rest.is_empty()
+            && first.is_name_expr()
+            && trailing_semicolon(item.into(), f.context().source()).is_some()
+        {
+            token(";").fmt(f)?;
+        }
+
+        Ok(())
     }
 
     fn is_suppressed(

crates/ruff_python_formatter/src/statement/stmt_aug_assign.rs

@@ -2,9 +2,11 @@ use ruff_formatter::write;
 use ruff_python_ast::StmtAugAssign;
 
 use crate::comments::{SourceComment, SuppressionKind};
+
 use crate::expression::maybe_parenthesize_expression;
 use crate::expression::parentheses::Parenthesize;
 use crate::prelude::*;
+use crate::statement::trailing_semicolon;
 use crate::{AsFormat, FormatNodeRule};
 
 #[derive(Default)]
@@ -28,7 +30,17 @@ impl FormatNodeRule<StmtAugAssign> for FormatStmtAugAssign {
                 space(),
                 maybe_parenthesize_expression(value, item, Parenthesize::IfBreaks)
             ]
-        )
+        )?;
+
+        if f.options().source_type().is_ipynb()
+            && f.context().node_level().is_last_top_level_statement()
+            && target.is_name_expr()
+            && trailing_semicolon(item.into(), f.context().source()).is_some()
+        {
+            token(";").fmt(f)?;
+        }
+
+        Ok(())
     }
 
     fn is_suppressed(

crates/ruff_python_formatter/src/statement/stmt_expr.rs

@@ -2,9 +2,11 @@ use ruff_python_ast as ast;
 use ruff_python_ast::{Expr, Operator, StmtExpr};
 
 use crate::comments::{SourceComment, SuppressionKind};
+
 use crate::expression::maybe_parenthesize_expression;
 use crate::expression::parentheses::Parenthesize;
 use crate::prelude::*;
+use crate::statement::trailing_semicolon;
 
 #[derive(Default)]
 pub struct FormatStmtExpr;
@@ -14,10 +16,19 @@ impl FormatNodeRule<StmtExpr> for FormatStmtExpr {
         let StmtExpr { value, .. } = item;
 
         if is_arithmetic_like(value) {
-            maybe_parenthesize_expression(value, item, Parenthesize::Optional).fmt(f)
+            maybe_parenthesize_expression(value, item, Parenthesize::Optional).fmt(f)?;
         } else {
-            value.format().fmt(f)
+            value.format().fmt(f)?;
         }
+
+        if f.options().source_type().is_ipynb()
+            && f.context().node_level().is_last_top_level_statement()
+            && trailing_semicolon(item.into(), f.context().source()).is_some()
+        {
+            token(";").fmt(f)?;
+        }
+
+        Ok(())
     }
 
     fn is_suppressed(

crates/ruff_python_formatter/src/statement/suite.rs

@@ -8,7 +8,7 @@ use ruff_text_size::{Ranged, TextRange};
 use crate::comments::{
     leading_comments, trailing_comments, Comments, LeadingDanglingTrailingComments,
 };
-use crate::context::{NodeLevel, WithNodeLevel};
+use crate::context::{NodeLevel, TopLevelStatementPosition, WithNodeLevel};
 use crate::expression::string::StringLayout;
 use crate::prelude::*;
 use crate::statement::stmt_expr::FormatStmtExpr;
@@ -49,8 +49,19 @@ impl Default for FormatSuite {
 
 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,
+            SuiteKind::TopLevel => NodeLevel::TopLevel(
+                iter.clone()
+                    .next()
+                    .map_or(TopLevelStatementPosition::Last, |_| {
+                        TopLevelStatementPosition::Other
+                    }),
+            ),
             SuiteKind::Function | SuiteKind::Class | SuiteKind::Other => {
                 NodeLevel::CompoundStatement
             }
@@ -62,11 +73,6 @@ impl FormatRule<Suite, PyFormatContext<'_>> for FormatSuite {
 
         let f = &mut WithNodeLevel::new(node_level, 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 first = match self.kind {
             SuiteKind::Other => {
@@ -165,6 +171,11 @@ impl FormatRule<Suite, PyFormatContext<'_>> for FormatSuite {
         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) {
@@ -351,7 +362,7 @@ impl FormatRule<Suite, PyFormatContext<'_>> for FormatSuite {
                     .map_or(preceding.end(), |comment| comment.slice().end());
 
                 match node_level {
-                    NodeLevel::TopLevel => match lines_after(end, source) {
+                    NodeLevel::TopLevel(_) => match lines_after(end, source) {
                         0 | 1 => hard_line_break().fmt(f)?,
                         2 => empty_line().fmt(f)?,
                         _ => match source_type {