Formatter: Add SourceType to context to enable special formatting for stub files (#6331)

**Summary** This adds the information whether we're in a .py python
source file or in a .pyi stub file to enable people working on #5822 and
related issues.

I'm not completely happy with `Default` for something that depends on
the input.

**Test Plan** None, this is currently unused, i'm leaving this to first
implementation of stub file specific formatting.

---------

Co-authored-by: Micha Reiser <micha@reiser.io>

crates/ruff_python_formatter/src/cli.rs

@@ -1,14 +1,14 @@
 #![allow(clippy::print_stdout)]
 
-use std::path::PathBuf;
+use std::path::{Path, PathBuf};
 
 use anyhow::{bail, Context, Result};
 use clap::{command, Parser, ValueEnum};
-use ruff_python_parser::lexer::lex;
-use ruff_python_parser::{parse_tokens, Mode};
 
 use ruff_formatter::SourceCode;
 use ruff_python_index::CommentRangesBuilder;
+use ruff_python_parser::lexer::lex;
+use ruff_python_parser::{parse_tokens, Mode};
 
 use crate::{format_node, PyFormatOptions};
 
@@ -37,7 +37,7 @@ pub struct Cli {
     pub print_comments: bool,
 }
 
-pub fn format_and_debug_print(input: &str, cli: &Cli) -> Result<String> {
+pub fn format_and_debug_print(input: &str, cli: &Cli, source_type: &Path) -> Result<String> {
     let mut tokens = Vec::new();
     let mut comment_ranges = CommentRangesBuilder::default();
 
@@ -57,13 +57,9 @@ pub fn format_and_debug_print(input: &str, cli: &Cli) -> Result<String> {
     let python_ast =
         parse_tokens(tokens, Mode::Module, "<filename>").context("Syntax error in input")?;
 
-    let formatted = format_node(
-        &python_ast,
-        &comment_ranges,
-        input,
-        PyFormatOptions::default(),
-    )
-    .context("Failed to format node")?;
+    let options = PyFormatOptions::from_extension(source_type);
+    let formatted = format_node(&python_ast, &comment_ranges, input, options)
+        .context("Failed to format node")?;
     if cli.print_ir {
         println!("{}", formatted.document().display(SourceCode::new(input)));
     }

crates/ruff_python_formatter/src/lib.rs

@@ -255,6 +255,7 @@ mod tests {
     use ruff_python_index::CommentRangesBuilder;
     use ruff_python_parser::lexer::lex;
     use ruff_python_parser::{parse_tokens, Mode};
+    use std::path::Path;
 
     /// Very basic test intentionally kept very similar to the CLI
     #[test]
@@ -321,15 +322,10 @@ with [
         let comment_ranges = comment_ranges.finish();
 
         // Parse the AST.
-        let python_ast = parse_tokens(tokens, Mode::Module, "<filename>").unwrap();
-
-        let formatted = format_node(
-            &python_ast,
-            &comment_ranges,
-            src,
-            PyFormatOptions::default(),
-        )
-        .unwrap();
+        let source_path = "code_inline.py";
+        let python_ast = parse_tokens(tokens, Mode::Module, source_path).unwrap();
+        let options = PyFormatOptions::from_extension(Path::new(source_path));
+        let formatted = format_node(&python_ast, &comment_ranges, src, options).unwrap();
 
         // Uncomment the `dbg` to print the IR.
         // Use `dbg_write!(f, []) instead of `write!(f, [])` in your formatting code to print some IR

crates/ruff_python_formatter/src/main.rs

@@ -1,4 +1,5 @@
 use std::io::{stdout, Read, Write};
+use std::path::Path;
 use std::{fs, io};
 
 use anyhow::{bail, Context, Result};
@@ -25,7 +26,8 @@ fn main() -> Result<()> {
             );
         }
         let input = read_from_stdin()?;
-        let formatted = format_and_debug_print(&input, &cli)?;
+        // It seems reasonable to give this a dummy name
+        let formatted = format_and_debug_print(&input, &cli, Path::new("stdin.py"))?;
         if cli.check {
             if formatted == input {
                 return Ok(());
@@ -37,7 +39,7 @@ fn main() -> Result<()> {
         for file in &cli.files {
             let input = fs::read_to_string(file)
                 .with_context(|| format!("Could not read {}: ", file.display()))?;
-            let formatted = format_and_debug_print(&input, &cli)?;
+            let formatted = format_and_debug_print(&input, &cli, file)?;
             match cli.emit {
                 Some(Emit::Stdout) => stdout().lock().write_all(formatted.as_bytes())?,
                 None | Some(Emit::Files) => {

crates/ruff_python_formatter/src/options.rs

@@ -1,5 +1,7 @@
 use ruff_formatter::printer::{LineEnding, PrinterOptions};
 use ruff_formatter::{FormatOptions, IndentStyle, LineWidth};
+use ruff_python_ast::PySourceType;
+use std::path::Path;
 
 #[derive(Clone, Debug)]
 #[cfg_attr(
@@ -8,6 +10,9 @@ use ruff_formatter::{FormatOptions, IndentStyle, LineWidth};
     serde(default)
 )]
 pub struct PyFormatOptions {
+    /// Whether we're in a `.py` file or `.pyi` file, which have different rules
+    source_type: PySourceType,
+
     /// Specifies the indent style:
     /// * Either a tab
     /// * or a specific amount of spaces
@@ -28,7 +33,31 @@ fn default_line_width() -> LineWidth {
     LineWidth::try_from(88).unwrap()
 }
 
+impl Default for PyFormatOptions {
+    fn default() -> Self {
+        Self {
+            source_type: PySourceType::default(),
+            indent_style: IndentStyle::Space(4),
+            line_width: LineWidth::try_from(88).unwrap(),
+            quote_style: QuoteStyle::default(),
+            magic_trailing_comma: MagicTrailingComma::default(),
+        }
+    }
+}
+
 impl PyFormatOptions {
+    /// Otherwise sets the defaults. Returns none if the extension is unknown
+    pub fn from_extension(path: &Path) -> Self {
+        Self::from_source_type(PySourceType::from(path))
+    }
+
+    pub fn from_source_type(source_type: PySourceType) -> Self {
+        Self {
+            source_type,
+            ..Self::default()
+        }
+    }
+
     pub fn magic_trailing_comma(&self) -> MagicTrailingComma {
         self.magic_trailing_comma
     }
@@ -42,17 +71,20 @@ impl PyFormatOptions {
         self
     }
 
-    pub fn with_magic_trailing_comma(&mut self, trailing_comma: MagicTrailingComma) -> &mut Self {
+    #[must_use]
+    pub fn with_magic_trailing_comma(mut self, trailing_comma: MagicTrailingComma) -> Self {
         self.magic_trailing_comma = trailing_comma;
         self
     }
 
-    pub fn with_indent_style(&mut self, indent_style: IndentStyle) -> &mut Self {
+    #[must_use]
+    pub fn with_indent_style(mut self, indent_style: IndentStyle) -> Self {
         self.indent_style = indent_style;
         self
     }
 
-    pub fn with_line_width(&mut self, line_width: LineWidth) -> &mut Self {
+    #[must_use]
+    pub fn with_line_width(mut self, line_width: LineWidth) -> Self {
         self.line_width = line_width;
         self
     }
@@ -77,17 +109,6 @@ impl FormatOptions for PyFormatOptions {
     }
 }
 
-impl Default for PyFormatOptions {
-    fn default() -> Self {
-        Self {
-            indent_style: IndentStyle::Space(4),
-            line_width: LineWidth::try_from(88).unwrap(),
-            quote_style: QuoteStyle::default(),
-            magic_trailing_comma: MagicTrailingComma::default(),
-        }
-    }
-}
-
 #[derive(Copy, Clone, Debug, Default, Eq, PartialEq)]
 #[cfg_attr(
     feature = "serde",