Abstract stylist to libcst style conversion

crates/ruff/src/autofix/codemods.rs

@@ -11,6 +11,23 @@ use ruff_python_ast::source_code::{Locator, Stylist};
 use crate::cst::helpers::compose_module_path;
 use crate::cst::matchers::match_statement;
 
+/// Glue code to make libcst codegen work with ruff's Stylist
+pub(crate) trait CodegenStylist<'a>: Codegen<'a> {
+    fn codegen_stylist(&self, stylist: &'a Stylist) -> String;
+}
+
+impl<'a, T: Codegen<'a>> CodegenStylist<'a> for T {
+    fn codegen_stylist(&self, stylist: &'a Stylist) -> String {
+        let mut state = CodegenState {
+            default_newline: stylist.line_ending().as_str(),
+            default_indent: stylist.indentation(),
+            ..Default::default()
+        };
+        self.codegen(&mut state);
+        state.to_string()
+    }
+}
+
 /// Given an import statement, remove any imports that are specified in the `imports` iterator.
 ///
 /// Returns `Ok(None)` if the statement is empty after removing the imports.
@@ -114,14 +131,7 @@ pub(crate) fn remove_imports<'a>(
         return Ok(None);
     }
 
-    let mut state = CodegenState {
-        default_newline: &stylist.line_ending(),
-        default_indent: stylist.indentation(),
-        ..CodegenState::default()
-    };
-    tree.codegen(&mut state);
-
-    Ok(Some(state.to_string()))
+    Ok(Some(tree.codegen_stylist(stylist)))
 }
 
 /// Given an import statement, remove any imports that are not specified in the `imports` slice.
@@ -200,11 +210,5 @@ pub(crate) fn retain_imports(
         }
     }
 
-    let mut state = CodegenState {
-        default_newline: &stylist.line_ending(),
-        default_indent: stylist.indentation(),
-        ..CodegenState::default()
-    };
-    tree.codegen(&mut state);
-    Ok(state.to_string())
+    Ok(tree.codegen_stylist(stylist))
 }

crates/ruff/src/importer/mod.rs

@@ -3,11 +3,12 @@
 use std::error::Error;
 
 use anyhow::Result;
-use libcst_native::{Codegen, CodegenState, ImportAlias, Name, NameOrAttribute};
+use libcst_native::{ImportAlias, Name, NameOrAttribute};
 use ruff_text_size::TextSize;
 use rustpython_parser::ast::{self, Ranged, Stmt, Suite};
 
 use crate::autofix;
+use crate::autofix::codemods::CodegenStylist;
 use ruff_diagnostics::Edit;
 use ruff_python_ast::imports::{AnyImport, Import, ImportFrom};
 use ruff_python_ast::source_code::{Locator, Stylist};
@@ -324,13 +325,10 @@ impl<'a> Importer<'a> {
             asname: None,
             comma: aliases.last().and_then(|alias| alias.comma.clone()),
         });
-        let mut state = CodegenState {
-            default_newline: &self.stylist.line_ending(),
-            default_indent: self.stylist.indentation(),
-            ..CodegenState::default()
-        };
-        statement.codegen(&mut state);
-        Ok(Edit::range_replacement(state.to_string(), stmt.range()))
+        Ok(Edit::range_replacement(
+            statement.codegen_stylist(self.stylist),
+            stmt.range(),
+        ))
     }
 
     /// Add a `TYPE_CHECKING` block to the given module.