Generate source code with detected line ending (#1487)

ruff_dev/src/round_trip.rs

@@ -22,7 +22,11 @@ pub fn main(cli: &Cli) -> Result<()> {
     let python_ast = parser::parse_program(&contents, &cli.file.to_string_lossy())?;
     let locator = SourceCodeLocator::new(&contents);
     let stylist = SourceCodeStyleDetector::from_contents(&contents, &locator);
-    let mut generator = SourceCodeGenerator::new(stylist.indentation(), stylist.quote());
+    let mut generator = SourceCodeGenerator::new(
+        stylist.indentation(),
+        stylist.quote(),
+        stylist.line_ending(),
+    );
     generator.unparse_suite(&python_ast);
     println!("{}", generator.generate()?);
     Ok(())

src/flake8_bugbear/plugins/assert_false.rs

@@ -47,8 +47,11 @@ pub fn assert_false(checker: &mut Checker, stmt: &Stmt, test: &Expr, msg: Option
 
     let mut check = Check::new(CheckKind::DoNotAssertFalse, Range::from_located(test));
     if checker.patch(check.kind.code()) {
-        let mut generator =
+        let mut generator = SourceCodeGenerator::new(
-            SourceCodeGenerator::new(checker.style.indentation(), checker.style.quote());
+            checker.style.indentation(),
+            checker.style.quote(),
+            checker.style.line_ending(),
+        );
         generator.unparse_stmt(&assertion_error(msg));
         if let Ok(content) = generator.generate() {
             check.amend(Fix::replacement(

src/flake8_bugbear/plugins/duplicate_exceptions.rs

@@ -54,8 +54,11 @@ fn duplicate_handler_exceptions<'a>(
                 Range::from_located(expr),
             );
             if checker.patch(check.kind.code()) {
-                let mut generator =
+                let mut generator = SourceCodeGenerator::new(
-                    SourceCodeGenerator::new(checker.style.indentation(), checker.style.quote());
+                    checker.style.indentation(),
+                    checker.style.quote(),
+                    checker.style.line_ending(),
+                );
                 if unique_elts.len() == 1 {
                     generator.unparse_expr(unique_elts[0], 0);
                 } else {

src/flake8_bugbear/plugins/getattr_with_constant.rs

@@ -46,8 +46,11 @@ pub fn getattr_with_constant(checker: &mut Checker, expr: &Expr, func: &Expr, ar
 
     let mut check = Check::new(CheckKind::GetAttrWithConstant, Range::from_located(expr));
     if checker.patch(check.kind.code()) {
-        let mut generator =
+        let mut generator = SourceCodeGenerator::new(
-            SourceCodeGenerator::new(checker.style.indentation(), checker.style.quote());
+            checker.style.indentation(),
+            checker.style.quote(),
+            checker.style.line_ending(),
+        );
         generator.unparse_expr(&attribute(obj, value), 0);
         if let Ok(content) = generator.generate() {
             check.amend(Fix::replacement(

src/flake8_bugbear/plugins/redundant_tuple_in_exception_handler.rs

@@ -23,8 +23,11 @@ pub fn redundant_tuple_in_exception_handler(checker: &mut Checker, handlers: &[E
             Range::from_located(type_),
         );
         if checker.patch(check.kind.code()) {
-            let mut generator =
+            let mut generator = SourceCodeGenerator::new(
-                SourceCodeGenerator::new(checker.style.indentation(), checker.style.quote());
+                checker.style.indentation(),
+                checker.style.quote(),
+                checker.style.line_ending(),
+            );
             generator.unparse_expr(elt, 0);
             if let Ok(content) = generator.generate() {
                 check.amend(Fix::replacement(

src/flake8_bugbear/plugins/setattr_with_constant.rs

@@ -34,7 +34,11 @@ fn assignment(
             type_comment: None,
         },
     );
-    let mut generator = SourceCodeGenerator::new(stylist.indentation(), stylist.quote());
+    let mut generator = SourceCodeGenerator::new(
+        stylist.indentation(),
+        stylist.quote(),
+        stylist.line_ending(),
+    );
     generator.unparse_stmt(&stmt);
     generator.generate().map_err(std::convert::Into::into)
 }

src/pycodestyle/plugins.rs

@@ -30,7 +30,11 @@ fn compare(
             comparators: comparators.to_vec(),
         },
     );
-    let mut generator = SourceCodeGenerator::new(stylist.indentation(), stylist.quote());
+    let mut generator = SourceCodeGenerator::new(
+        stylist.indentation(),
+        stylist.quote(),
+        stylist.line_ending(),
+    );
     generator.unparse_expr(&cmp, 0);
     generator.generate().ok()
 }
@@ -302,7 +306,11 @@ fn function(
             type_comment: None,
         },
     );
-    let mut generator = SourceCodeGenerator::new(stylist.indentation(), stylist.quote());
+    let mut generator = SourceCodeGenerator::new(
+        stylist.indentation(),
+        stylist.quote(),
+        stylist.line_ending(),
+    );
     generator.unparse_stmt(&func);
     Ok(generator.generate()?)
 }

src/pyupgrade/plugins/convert_named_tuple_functional_to_class.rs

@@ -166,7 +166,11 @@ fn convert_to_class(
     base_class: &ExprKind,
     stylist: &SourceCodeStyleDetector,
 ) -> Result<Fix> {
-    let mut generator = SourceCodeGenerator::new(stylist.indentation(), stylist.quote());
+    let mut generator = SourceCodeGenerator::new(
+        stylist.indentation(),
+        stylist.quote(),
+        stylist.line_ending(),
+    );
     generator.unparse_stmt(&create_class_def_stmt(typename, body, base_class));
     let content = generator.generate()?;
     Ok(Fix::replacement(

src/pyupgrade/plugins/convert_typed_dict_functional_to_class.rs

@@ -199,7 +199,11 @@ fn convert_to_class(
     base_class: &ExprKind,
     stylist: &SourceCodeStyleDetector,
 ) -> Result<Fix> {
-    let mut generator = SourceCodeGenerator::new(stylist.indentation(), stylist.quote());
+    let mut generator = SourceCodeGenerator::new(
+        stylist.indentation(),
+        stylist.quote(),
+        stylist.line_ending(),
+    );
     generator.unparse_stmt(&create_class_def_stmt(
         class_name,
         body,

src/pyupgrade/plugins/remove_six_compat.rs

@@ -138,7 +138,11 @@ fn replace_by_expr_kind(
 ) -> Result<Check> {
     let mut check = Check::new(CheckKind::RemoveSixCompat, Range::from_located(expr));
     if patch {
-        let mut generator = SourceCodeGenerator::new(stylist.indentation(), stylist.quote());
+        let mut generator = SourceCodeGenerator::new(
+            stylist.indentation(),
+            stylist.quote(),
+            stylist.line_ending(),
+        );
         generator.unparse_expr(&create_expr(node), 0);
         let content = generator.generate()?;
         check.amend(Fix::replacement(
@@ -158,7 +162,11 @@ fn replace_by_stmt_kind(
 ) -> Result<Check> {
     let mut check = Check::new(CheckKind::RemoveSixCompat, Range::from_located(expr));
     if patch {
-        let mut generator = SourceCodeGenerator::new(stylist.indentation(), stylist.quote());
+        let mut generator = SourceCodeGenerator::new(
+            stylist.indentation(),
+            stylist.quote(),
+            stylist.line_ending(),
+        );
         generator.unparse_stmt(&create_stmt(node));
         let content = generator.generate()?;
         check.amend(Fix::replacement(

src/pyupgrade/plugins/use_pep604_annotation.rs

@@ -65,8 +65,11 @@ pub fn use_pep604_annotation(checker: &mut Checker, expr: &Expr, value: &Expr, s
     if checker.match_typing_call_path(&call_path, "Optional") {
         let mut check = Check::new(CheckKind::UsePEP604Annotation, Range::from_located(expr));
         if checker.patch(check.kind.code()) {
-            let mut generator =
+            let mut generator = SourceCodeGenerator::new(
-                SourceCodeGenerator::new(checker.style.indentation(), checker.style.quote());
+                checker.style.indentation(),
+                checker.style.quote(),
+                checker.style.line_ending(),
+            );
             generator.unparse_expr(&optional(slice), 0);
             if let Ok(content) = generator.generate() {
                 check.amend(Fix::replacement(
@@ -88,6 +91,7 @@ pub fn use_pep604_annotation(checker: &mut Checker, expr: &Expr, value: &Expr, s
                     let mut generator = SourceCodeGenerator::new(
                         checker.style.indentation(),
                         checker.style.quote(),
+                        checker.style.line_ending(),
                     );
                     generator.unparse_expr(&union(elts), 0);
                     if let Ok(content) = generator.generate() {
@@ -103,6 +107,7 @@ pub fn use_pep604_annotation(checker: &mut Checker, expr: &Expr, value: &Expr, s
                     let mut generator = SourceCodeGenerator::new(
                         checker.style.indentation(),
                         checker.style.quote(),
+                        checker.style.line_ending(),
                     );
                     generator.unparse_expr(slice, 0);
                     if let Ok(content) = generator.generate() {

src/source_code_generator.rs

@@ -11,7 +11,7 @@ use rustpython_parser::ast::{
     Operator, Stmt, StmtKind,
 };
 
-use crate::source_code_style::{Indentation, Quote};
+use crate::source_code_style::{Indentation, LineEnding, Quote};
 use crate::vendor::{bytes, str};
 
 mod precedence {
@@ -37,6 +37,8 @@ pub struct SourceCodeGenerator<'a> {
     indent: &'a Indentation,
     /// The quote style to use for string literals.
     quote: &'a Quote,
+    /// The line ending to use.
+    line_ending: &'a LineEnding,
     buffer: Vec<u8>,
     indent_depth: usize,
     num_newlines: usize,
@@ -44,11 +46,12 @@ pub struct SourceCodeGenerator<'a> {
 }
 
 impl<'a> SourceCodeGenerator<'a> {
-    pub fn new(indent: &'a Indentation, quote: &'a Quote) -> Self {
+    pub fn new(indent: &'a Indentation, quote: &'a Quote, line_ending: &'a LineEnding) -> Self {
         SourceCodeGenerator {
             // Style preferences.
             indent,
             quote,
+            line_ending,
             // Internal state.
             buffer: vec![],
             indent_depth: 0,
@@ -84,7 +87,7 @@ impl<'a> SourceCodeGenerator<'a> {
     fn p(&mut self, s: &str) {
         if self.num_newlines > 0 {
             for _ in 0..self.num_newlines {
-                self.buffer.extend("\n".as_bytes());
+                self.buffer.extend(self.line_ending.as_bytes());
             }
             self.num_newlines = 0;
         }
@@ -944,7 +947,7 @@ impl<'a> SourceCodeGenerator<'a> {
     }
 
     fn unparse_formatted<U>(&mut self, val: &Expr<U>, conversion: usize, spec: Option<&Expr<U>>) {
-        let mut generator = SourceCodeGenerator::new(self.indent, self.quote);
+        let mut generator = SourceCodeGenerator::new(self.indent, self.quote, self.line_ending);
         generator.unparse_expr(val, precedence::TEST + 1);
         let brace = if generator.buffer.starts_with("{".as_bytes()) {
             // put a space to avoid escaping the bracket
@@ -1000,7 +1003,7 @@ impl<'a> SourceCodeGenerator<'a> {
             self.unparse_fstring_body(values, is_spec);
         } else {
             self.p("f");
-            let mut generator = SourceCodeGenerator::new(self.indent, self.quote);
+            let mut generator = SourceCodeGenerator::new(self.indent, self.quote, self.line_ending);
             generator.unparse_fstring_body(values, is_spec);
             let body = std::str::from_utf8(&generator.buffer).unwrap();
             self.p(&format!("{}", str::repr(body, self.quote.into())));
@@ -1031,22 +1034,28 @@ mod tests {
     use rustpython_parser::parser;
 
     use crate::source_code_generator::SourceCodeGenerator;
-    use crate::source_code_style::{Indentation, Quote};
+    use crate::source_code_style::{Indentation, LineEnding, Quote};
 
     fn round_trip(contents: &str) -> Result<String> {
         let indentation = Indentation::default();
         let quote = Quote::default();
+        let line_ending = LineEnding::default();
         let program = parser::parse_program(contents, "<filename>")?;
         let stmt = program.first().unwrap();
-        let mut generator = SourceCodeGenerator::new(&indentation, &quote);
+        let mut generator = SourceCodeGenerator::new(&indentation, &quote, &line_ending);
         generator.unparse_stmt(stmt);
         generator.generate().map_err(std::convert::Into::into)
     }
 
-    fn round_trip_with(indentation: &Indentation, quote: &Quote, contents: &str) -> Result<String> {
+    fn round_trip_with(
+        indentation: &Indentation,
+        quote: &Quote,
+        line_ending: &LineEnding,
+        contents: &str,
+    ) -> Result<String> {
         let program = parser::parse_program(contents, "<filename>")?;
         let stmt = program.first().unwrap();
-        let mut generator = SourceCodeGenerator::new(indentation, quote);
+        let mut generator = SourceCodeGenerator::new(indentation, quote, line_ending);
         generator.unparse_stmt(stmt);
         generator.generate().map_err(std::convert::Into::into)
     }
@@ -1087,19 +1096,39 @@ if True:
     #[test]
     fn set_quote() -> Result<()> {
         assert_eq!(
-            round_trip_with(&Indentation::default(), &Quote::Double, r#""hello""#)?,
+            round_trip_with(
+                &Indentation::default(),
+                &Quote::Double,
+                &LineEnding::default(),
+                r#""hello""#
+            )?,
             r#""hello""#
         );
         assert_eq!(
-            round_trip_with(&Indentation::default(), &Quote::Single, r#""hello""#)?,
+            round_trip_with(
+                &Indentation::default(),
+                &Quote::Single,
+                &LineEnding::default(),
+                r#""hello""#
+            )?,
             r#"'hello'"#
         );
         assert_eq!(
-            round_trip_with(&Indentation::default(), &Quote::Double, r#"'hello'"#)?,
+            round_trip_with(
+                &Indentation::default(),
+                &Quote::Double,
+                &LineEnding::default(),
+                r#"'hello'"#
+            )?,
             r#""hello""#
         );
         assert_eq!(
-            round_trip_with(&Indentation::default(), &Quote::Single, r#"'hello'"#)?,
+            round_trip_with(
+                &Indentation::default(),
+                &Quote::Single,
+                &LineEnding::default(),
+                r#"'hello'"#
+            )?,
             r#"'hello'"#
         );
         Ok(())
@@ -1111,6 +1140,7 @@ if True:
             round_trip_with(
                 &Indentation::new("    ".to_string()),
                 &Quote::default(),
+                &LineEnding::default(),
                 r#"
 if True:
   pass
@@ -1127,6 +1157,7 @@ if True:
             round_trip_with(
                 &Indentation::new("  ".to_string()),
                 &Quote::default(),
+                &LineEnding::default(),
                 r#"
 if True:
   pass
@@ -1143,6 +1174,7 @@ if True:
             round_trip_with(
                 &Indentation::new("\t".to_string()),
                 &Quote::default(),
+                &LineEnding::default(),
                 r#"
 if True:
   pass
@@ -1158,4 +1190,39 @@ if True:
 
         Ok(())
     }
+
+    #[test]
+    fn set_line_ending() -> Result<()> {
+        assert_eq!(
+            round_trip_with(
+                &Indentation::default(),
+                &Quote::default(),
+                &LineEnding::Lf,
+                "if True:\n    print(42)",
+            )?,
+            "if True:\n    print(42)",
+        );
+
+        assert_eq!(
+            round_trip_with(
+                &Indentation::default(),
+                &Quote::default(),
+                &LineEnding::CrLf,
+                "if True:\n    print(42)",
+            )?,
+            "if True:\r\n    print(42)",
+        );
+
+        assert_eq!(
+            round_trip_with(
+                &Indentation::default(),
+                &Quote::default(),
+                &LineEnding::Cr,
+                "if True:\n    print(42)",
+            )?,
+            "if True:\r    print(42)",
+        );
+
+        Ok(())
+    }
 }