Formatter parentheses support for IpyEscapeCommand

crates/ruff_cli/resources/test/fixtures/unformatted.ipynb

@@ -11,6 +11,41 @@
     "maths = (numpy.arange(100)**2).sum()\n",
     "stats= numpy.asarray([1,2,3,4]).median()"
    ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "83a0b1b8",
+   "metadata": {},
+   "source": [
+    "A markdown cell"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "ae12f012",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# A cell with IPython escape command\n",
+    "def some_function(foo, bar):\n",
+    "    pass\n",
+    "%matplotlib inline"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "10f3bbf9",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "foo = %pwd\n",
+    "def some_function(foo,bar,):\n",
+    "    # Another cell with IPython escape command\n",
+    "    foo = %pwd\n",
+    "    print(foo)"
+   ]
   }
  ],
  "metadata": {

crates/ruff_cli/tests/format.rs

@@ -419,13 +419,30 @@ fn test_diff() {
         ----- stdout -----
         --- resources/test/fixtures/unformatted.ipynb
         +++ resources/test/fixtures/unformatted.ipynb
-        @@ -1,3 +1,4 @@
+        @@ -1,12 +1,20 @@
          import numpy
         -maths = (numpy.arange(100)**2).sum()
         -stats= numpy.asarray([1,2,3,4]).median()
         +
         +maths = (numpy.arange(100) ** 2).sum()
         +stats = numpy.asarray([1, 2, 3, 4]).median()
+         # A cell with IPython escape command
+         def some_function(foo, bar):
+             pass
+        +
+        +
+         %matplotlib inline
+         foo = %pwd
+        -def some_function(foo,bar,):
+        +
+        +
+        +def some_function(
+        +    foo,
+        +    bar,
+        +):
+             # Another cell with IPython escape command
+             foo = %pwd
+             print(foo)
         --- resources/test/fixtures/unformatted.py
         +++ resources/test/fixtures/unformatted.py
         @@ -1,3 +1,3 @@

crates/ruff_python_formatter/src/cli.rs

@@ -6,8 +6,9 @@ use anyhow::{format_err, Context, Result};
 use clap::{command, Parser, ValueEnum};
 
 use ruff_formatter::SourceCode;
+use ruff_python_ast::PySourceType;
 use ruff_python_index::tokens_and_ranges;
-use ruff_python_parser::{parse_ok_tokens, Mode};
+use ruff_python_parser::{parse_ok_tokens, AsMode};
 use ruff_text_size::Ranged;
 
 use crate::comments::collect_comments;
@@ -38,15 +39,16 @@ pub struct Cli {
     pub print_comments: bool,
 }
 
-pub fn format_and_debug_print(source: &str, cli: &Cli, source_type: &Path) -> Result<String> {
-    let (tokens, comment_ranges) = tokens_and_ranges(source)
+pub fn format_and_debug_print(source: &str, cli: &Cli, source_path: &Path) -> Result<String> {
+    let source_type = PySourceType::from(source_path);
+    let (tokens, comment_ranges) = tokens_and_ranges(source, source_type)
         .map_err(|err| format_err!("Source contains syntax errors {err:?}"))?;
 
     // Parse the AST.
-    let module = parse_ok_tokens(tokens, source, Mode::Module, "<filename>")
+    let module = parse_ok_tokens(tokens, source, source_type.as_mode(), "<filename>")
         .context("Syntax error in input")?;
 
-    let options = PyFormatOptions::from_extension(source_type);
+    let options = PyFormatOptions::from_extension(source_path);
 
     let source_code = SourceCode::new(source);
     let formatted = format_module_ast(&module, &comment_ranges, source, options)

crates/ruff_python_formatter/src/comments/mod.rs

@@ -548,10 +548,10 @@ mod tests {
     use insta::assert_debug_snapshot;
 
     use ruff_formatter::SourceCode;
-    use ruff_python_ast::Mod;
+    use ruff_python_ast::{Mod, PySourceType};
     use ruff_python_index::tokens_and_ranges;
 
-    use ruff_python_parser::{parse_ok_tokens, Mode};
+    use ruff_python_parser::{parse_ok_tokens, AsMode};
     use ruff_python_trivia::CommentRanges;
 
     use crate::comments::Comments;
@@ -565,9 +565,10 @@ mod tests {
     impl<'a> CommentsTestCase<'a> {
         fn from_code(source: &'a str) -> Self {
             let source_code = SourceCode::new(source);
+            let source_type = PySourceType::Python;
             let (tokens, comment_ranges) =
-                tokens_and_ranges(source).expect("Expect source to be valid Python");
-            let parsed = parse_ok_tokens(tokens, source, Mode::Module, "test.py")
+                tokens_and_ranges(source, source_type).expect("Expect source to be valid Python");
+            let parsed = parse_ok_tokens(tokens, source, source_type.as_mode(), "test.py")
                 .expect("Expect source to be valid Python");
 
             CommentsTestCase {

crates/ruff_python_formatter/src/expression/expr_ipy_escape_command.rs

@@ -1,6 +1,7 @@
 use ruff_python_ast::ExprIpyEscapeCommand;
 use ruff_text_size::Ranged;
 
+use crate::expression::parentheses::{NeedsParentheses, OptionalParentheses};
 use crate::prelude::*;
 
 #[derive(Default)]
@@ -11,3 +12,13 @@ impl FormatNodeRule<ExprIpyEscapeCommand> for FormatExprIpyEscapeCommand {
         source_text_slice(item.range()).fmt(f)
     }
 }
+
+impl NeedsParentheses for ExprIpyEscapeCommand {
+    fn needs_parentheses(
+        &self,
+        _parent: ruff_python_ast::AnyNodeRef,
+        _context: &PyFormatContext,
+    ) -> OptionalParentheses {
+        OptionalParentheses::Never
+    }
+}

crates/ruff_python_formatter/src/expression/mod.rs

@@ -476,7 +476,7 @@ impl NeedsParentheses for Expr {
             Expr::List(expr) => expr.needs_parentheses(parent, context),
             Expr::Tuple(expr) => expr.needs_parentheses(parent, context),
             Expr::Slice(expr) => expr.needs_parentheses(parent, context),
-            Expr::IpyEscapeCommand(_) => todo!(),
+            Expr::IpyEscapeCommand(expr) => expr.needs_parentheses(parent, context),
         }
     }
 }
@@ -718,8 +718,8 @@ impl<'input> CanOmitOptionalParenthesesVisitor<'input> {
             | Expr::Constant(_)
             | Expr::Starred(_)
             | Expr::Name(_)
-            | Expr::Slice(_) => {}
-            Expr::IpyEscapeCommand(_) => todo!(),
+            | Expr::Slice(_)
+            | Expr::IpyEscapeCommand(_) => {}
         };
 
         walk_expr(self, expr);

crates/ruff_python_formatter/src/lib.rs

@@ -7,7 +7,7 @@ use ruff_python_ast::AstNode;
 use ruff_python_ast::Mod;
 use ruff_python_index::tokens_and_ranges;
 use ruff_python_parser::lexer::LexicalError;
-use ruff_python_parser::{parse_ok_tokens, Mode, ParseError};
+use ruff_python_parser::{parse_ok_tokens, AsMode, ParseError};
 use ruff_python_trivia::CommentRanges;
 use ruff_source_file::Locator;
 
@@ -130,8 +130,9 @@ pub fn format_module_source(
     source: &str,
     options: PyFormatOptions,
 ) -> Result<Printed, FormatModuleError> {
-    let (tokens, comment_ranges) = tokens_and_ranges(source)?;
-    let module = parse_ok_tokens(tokens, source, Mode::Module, "<filename>")?;
+    let source_type = options.source_type();
+    let (tokens, comment_ranges) = tokens_and_ranges(source, source_type)?;
+    let module = parse_ok_tokens(tokens, source, source_type.as_mode(), "<filename>")?;
     let formatted = format_module_ast(&module, &comment_ranges, source, options)?;
     Ok(formatted.print()?)
 }
@@ -172,9 +173,10 @@ mod tests {
     use anyhow::Result;
     use insta::assert_snapshot;
 
+    use ruff_python_ast::PySourceType;
     use ruff_python_index::tokens_and_ranges;
 
-    use ruff_python_parser::{parse_ok_tokens, Mode};
+    use ruff_python_parser::{parse_ok_tokens, AsMode};
 
     use crate::{format_module_ast, format_module_source, PyFormatOptions};
 
@@ -213,11 +215,12 @@ def main() -> None:
         ]
 
 "#;
-        let (tokens, comment_ranges) = tokens_and_ranges(source).unwrap();
+        let source_type = PySourceType::Python;
+        let (tokens, comment_ranges) = tokens_and_ranges(source, source_type).unwrap();
 
         // Parse the AST.
         let source_path = "code_inline.py";
-        let module = parse_ok_tokens(tokens, source, Mode::Module, source_path).unwrap();
+        let module = parse_ok_tokens(tokens, source, source_type.as_mode(), source_path).unwrap();
         let options = PyFormatOptions::from_extension(Path::new(source_path));
         let formatted = format_module_ast(&module, &comment_ranges, source, options).unwrap();
 

crates/ruff_python_index/src/comment_ranges.rs

@@ -1,7 +1,8 @@
 use std::fmt::Debug;
 
+use ruff_python_ast::PySourceType;
 use ruff_python_parser::lexer::{lex, LexicalError};
-use ruff_python_parser::{Mode, Tok};
+use ruff_python_parser::{AsMode, Tok};
 use ruff_python_trivia::CommentRanges;
 use ruff_text_size::TextRange;
 
@@ -25,11 +26,12 @@ impl CommentRangesBuilder {
 /// Helper method to lex and extract comment ranges
 pub fn tokens_and_ranges(
     source: &str,
+    source_type: PySourceType,
 ) -> Result<(Vec<(Tok, TextRange)>, CommentRanges), LexicalError> {
     let mut tokens = Vec::new();
     let mut comment_ranges = CommentRangesBuilder::default();
 
-    for result in lex(source, Mode::Module) {
+    for result in lex(source, source_type.as_mode()) {
         let (token, range) = result?;
 
         comment_ranges.visit_token(&token, range);