Use a Rope to power fixer (#584)

src/autofix/fixer.rs

@@ -1,10 +1,14 @@
+use std::borrow::Cow;
 use std::collections::BTreeSet;
 
 use itertools::Itertools;
+use ropey::RopeBuilder;
 use rustpython_parser::ast::Location;
 
+use crate::ast::types::Range;
 use crate::autofix::{Fix, Patch};
 use crate::checks::Check;
+use crate::source_code_locator::SourceCodeLocator;
 
 // TODO(charlie): The model here is awkward because `Apply` is only relevant at
 // higher levels in the execution flow.
@@ -36,7 +40,7 @@ impl From<bool> for Mode {
 }
 
 /// Auto-fix errors in a file, and write the fixed source code to disk.
-pub fn fix_file(checks: &mut [Check], contents: &str) -> Option<String> {
+pub fn fix_file<'a>(checks: &mut [Check], contents: &'a str) -> Option<Cow<'a, str>> {
     if checks.iter().all(|check| check.fix.is_none()) {
         return None;
     }
@@ -48,12 +52,12 @@ pub fn fix_file(checks: &mut [Check], contents: &str) -> Option<String> {
 }
 
 /// Apply a series of fixes.
-fn apply_fixes<'a>(fixes: impl Iterator<Item = &'a mut Fix>, contents: &str) -> String {
-    let lines: Vec<&str> = contents.lines().collect();
+fn apply_fixes<'a>(fixes: impl Iterator<Item = &'a mut Fix>, contents: &str) -> Cow<'_, str> {
+    let locator = SourceCodeLocator::new(contents);
 
-    let mut output: String = Default::default();
-    let mut last_pos: Location = Default::default();
-    let mut applied: BTreeSet<&Patch> = Default::default();
+    let mut output = RopeBuilder::new();
+    let mut last_pos: Location = Location::new(1, 0);
+    let mut applied: BTreeSet<&Patch> = BTreeSet::new();
 
     for fix in fixes.sorted_by_key(|fix| fix.patch.location) {
         // If we already applied an identical fix as part of another correction, skip
@@ -69,44 +73,27 @@ fn apply_fixes<'a>(fixes: impl Iterator<Item = &'a mut Fix>, contents: &str) ->
             continue;
         }
 
-        if fix.patch.location.row() > last_pos.row() {
-            if last_pos.row() > 0 || last_pos.column() > 0 {
-                output.push_str(&lines[last_pos.row() - 1][last_pos.column()..]);
-                output.push('\n');
-            }
-            for line in &lines[last_pos.row()..fix.patch.location.row() - 1] {
-                output.push_str(line);
-                output.push('\n');
-            }
-            output.push_str(&lines[fix.patch.location.row() - 1][..fix.patch.location.column()]);
-            output.push_str(&fix.patch.content);
-        } else {
-            output.push_str(
-                &lines[last_pos.row() - 1][last_pos.column()..fix.patch.location.column()],
-            );
-            output.push_str(&fix.patch.content);
-        }
-        last_pos = fix.patch.end_location;
+        // Add all contents from `last_pos` to `fix.patch.location`.
+        let slice = locator.slice_source_code_range(&Range {
+            location: last_pos,
+            end_location: fix.patch.location,
+        });
+        output.append(&slice);
 
+        // Add the patch itself.
+        output.append(&fix.patch.content);
+
+        // Track that the fix was applied.
+        last_pos = fix.patch.end_location;
         applied.insert(&fix.patch);
         fix.applied = true;
     }
 
-    if last_pos.row() > 0
-        && (last_pos.row() - 1) < lines.len()
-        && (last_pos.row() > 0 || last_pos.column() > 0)
-    {
-        output.push_str(&lines[last_pos.row() - 1][last_pos.column()..]);
-        output.push('\n');
-    }
-    if last_pos.row() < lines.len() {
-        for line in &lines[last_pos.row()..] {
-            output.push_str(line);
-            output.push('\n');
-        }
-    }
+    // Add the remaining content.
+    let slice = locator.slice_source_code_at(&last_pos);
+    output.append(&slice);
 
-    output
+    Cow::from(output.finish())
 }
 
 #[cfg(test)]

src/cst/matchers.rs

@@ -1,18 +1,8 @@
 use anyhow::Result;
 use libcst_native::Module;
-use rustpython_ast::Located;
 
-use crate::ast::types::Range;
-use crate::source_code_locator::SourceCodeLocator;
-
-pub fn match_tree<'a, T>(
-    locator: &'a SourceCodeLocator,
-    located: &'a Located<T>,
-) -> Result<Module<'a>> {
-    match libcst_native::parse_module(
-        locator.slice_source_code_range(&Range::from_located(located)),
-        None,
-    ) {
+pub fn match_module(module_text: &str) -> Result<Module> {
+    match libcst_native::parse_module(module_text, None) {
         Ok(module) => Ok(module),
         Err(_) => return Err(anyhow::anyhow!("Failed to extract CST from source.")),
     }

src/docstrings/helpers.rs

@@ -25,12 +25,15 @@ pub fn leading_space(line: &str) -> String {
 }
 
 /// Extract the leading indentation from a docstring.
-pub fn indentation<'a>(checker: &'a Checker, docstring: &Expr) -> &'a str {
+pub fn indentation<'a>(checker: &'a Checker, docstring: &Expr) -> String {
     let range = Range::from_located(docstring);
-    checker.locator.slice_source_code_range(&Range {
-        location: Location::new(range.location.row(), 0),
-        end_location: Location::new(range.location.row(), range.location.column()),
-    })
+    checker
+        .locator
+        .slice_source_code_range(&Range {
+            location: Location::new(range.location.row(), 0),
+            end_location: Location::new(range.location.row(), range.location.column()),
+        })
+        .to_string()
 }
 
 /// Replace any non-whitespace characters from an indentation string.

src/flake8_comprehensions/fixes.rs

@@ -6,8 +6,9 @@ use libcst_native::{
     SmallStatement, Statement, Tuple,
 };
 
+use crate::ast::types::Range;
 use crate::autofix::Fix;
-use crate::cst::matchers::match_tree;
+use crate::cst::matchers::match_module;
 use crate::source_code_locator::SourceCodeLocator;
 
 fn match_expr<'a, 'b>(module: &'a mut Module<'b>) -> Result<&'a mut Expr<'b>> {
@@ -44,7 +45,8 @@ pub fn fix_unnecessary_generator_list(
     expr: &rustpython_ast::Expr,
 ) -> Result<Fix> {
     // Expr(Call(GeneratorExp)))) -> Expr(ListComp)))
-    let mut tree = match_tree(locator, expr)?;
+    let module_text = locator.slice_source_code_range(&Range::from_located(expr));
+    let mut tree = match_module(&module_text)?;
     let mut body = match_expr(&mut tree)?;
     let call = match_call(body)?;
     let arg = match_arg(call)?;
@@ -86,7 +88,8 @@ pub fn fix_unnecessary_generator_set(
     expr: &rustpython_ast::Expr,
 ) -> Result<Fix> {
     // Expr(Call(GeneratorExp)))) -> Expr(SetComp)))
-    let mut tree = match_tree(locator, expr)?;
+    let module_text = locator.slice_source_code_range(&Range::from_located(expr));
+    let mut tree = match_module(&module_text)?;
     let mut body = match_expr(&mut tree)?;
     let call = match_call(body)?;
     let arg = match_arg(call)?;
@@ -128,7 +131,8 @@ pub fn fix_unnecessary_generator_dict(
     locator: &SourceCodeLocator,
     expr: &rustpython_ast::Expr,
 ) -> Result<Fix> {
-    let mut tree = match_tree(locator, expr)?;
+    let module_text = locator.slice_source_code_range(&Range::from_located(expr));
+    let mut tree = match_module(&module_text)?;
     let mut body = match_expr(&mut tree)?;
     let call = match_call(body)?;
     let arg = match_arg(call)?;
@@ -194,7 +198,8 @@ pub fn fix_unnecessary_list_comprehension_set(
 ) -> Result<Fix> {
     // Expr(Call(ListComp)))) ->
     // Expr(SetComp)))
-    let mut tree = match_tree(locator, expr)?;
+    let module_text = locator.slice_source_code_range(&Range::from_located(expr));
+    let mut tree = match_module(&module_text)?;
     let mut body = match_expr(&mut tree)?;
     let call = match_call(body)?;
     let arg = match_arg(call)?;
@@ -234,7 +239,8 @@ pub fn fix_unnecessary_literal_set(
     expr: &rustpython_ast::Expr,
 ) -> Result<Fix> {
     // Expr(Call(List|Tuple)))) -> Expr(Set)))
-    let mut tree = match_tree(locator, expr)?;
+    let module_text = locator.slice_source_code_range(&Range::from_located(expr));
+    let mut tree = match_module(&module_text)?;
     let mut body = match_expr(&mut tree)?;
     let mut call = match_call(body)?;
     let arg = match_arg(call)?;
@@ -281,7 +287,8 @@ pub fn fix_unnecessary_literal_dict(
     expr: &rustpython_ast::Expr,
 ) -> Result<Fix> {
     // Expr(Call(List|Tuple)))) -> Expr(Dict)))
-    let mut tree = match_tree(locator, expr)?;
+    let module_text = locator.slice_source_code_range(&Range::from_located(expr));
+    let mut tree = match_module(&module_text)?;
     let mut body = match_expr(&mut tree)?;
     let call = match_call(body)?;
     let arg = match_arg(call)?;
@@ -351,8 +358,9 @@ pub fn fix_unnecessary_collection_call(
     locator: &SourceCodeLocator,
     expr: &rustpython_ast::Expr,
 ) -> Result<Fix> {
-    // Expr(Call("list" | "tuple" | "dict")))) -> Expr(List|Tuple|Dict)))
-    let mut tree = match_tree(locator, expr)?;
+    // Expr(Call("list" | "tuple" | "dict")))) -> Expr(List|Tuple|Dict)
+    let module_text = locator.slice_source_code_range(&Range::from_located(expr));
+    let mut tree = match_module(&module_text)?;
     let mut body = match_expr(&mut tree)?;
     let call = match_call(body)?;
     let name = if let Expression::Name(name) = &call.func.as_ref() {
@@ -463,7 +471,8 @@ pub fn fix_unnecessary_literal_within_tuple_call(
     locator: &SourceCodeLocator,
     expr: &rustpython_ast::Expr,
 ) -> Result<Fix> {
-    let mut tree = match_tree(locator, expr)?;
+    let module_text = locator.slice_source_code_range(&Range::from_located(expr));
+    let mut tree = match_module(&module_text)?;
     let mut body = match_expr(&mut tree)?;
     let call = match_call(body)?;
     let arg = match_arg(call)?;
@@ -518,7 +527,8 @@ pub fn fix_unnecessary_literal_within_list_call(
     locator: &SourceCodeLocator,
     expr: &rustpython_ast::Expr,
 ) -> Result<Fix> {
-    let mut tree = match_tree(locator, expr)?;
+    let module_text = locator.slice_source_code_range(&Range::from_located(expr));
+    let mut tree = match_module(&module_text)?;
     let mut body = match_expr(&mut tree)?;
     let call = match_call(body)?;
     let arg = match_arg(call)?;
@@ -576,7 +586,8 @@ pub fn fix_unnecessary_list_call(
     expr: &rustpython_ast::Expr,
 ) -> Result<Fix> {
     // Expr(Call(List|Tuple)))) -> Expr(List|Tuple)))
-    let mut tree = match_tree(locator, expr)?;
+    let module_text = locator.slice_source_code_range(&Range::from_located(expr));
+    let mut tree = match_module(&module_text)?;
     let mut body = match_expr(&mut tree)?;
     let call = match_call(body)?;
     let arg = match_arg(call)?;
@@ -598,7 +609,8 @@ pub fn fix_unnecessary_comprehension(
     locator: &SourceCodeLocator,
     expr: &rustpython_ast::Expr,
 ) -> Result<Fix> {
-    let mut tree = match_tree(locator, expr)?;
+    let module_text = locator.slice_source_code_range(&Range::from_located(expr));
+    let mut tree = match_module(&module_text)?;
     let mut body = match_expr(&mut tree)?;
 
     match &body.value {

src/linter.rs

@@ -131,11 +131,10 @@ pub fn lint_stdin(
 
     // Apply autofix, write results to stdout.
     if matches!(autofix, fixer::Mode::Apply) {
-        let output = match fix_file(&mut checks, stdin) {
-            None => stdin.to_string(),
-            Some(content) => content,
-        };
-        io::stdout().write_all(output.as_bytes())?;
+        match fix_file(&mut checks, stdin) {
+            None => io::stdout().write_all(stdin.as_bytes()),
+            Some(contents) => io::stdout().write_all(contents.as_bytes()),
+        }?;
     }
 
     // Convert to messages.
@@ -176,7 +175,7 @@ pub fn lint_path(
     // Apply autofix.
     if matches!(autofix, fixer::Mode::Apply) {
         if let Some(fixed_contents) = fix_file(&mut checks, &contents) {
-            write(path, fixed_contents)?;
+            write(path, fixed_contents.as_ref())?;
         }
     };
 

src/pycodestyle/checks.rs

@@ -271,7 +271,7 @@ pub fn invalid_escape_sequence(
     });
 
     // Determine whether the string is single- or triple-quoted.
-    let quote = extract_quote(text);
+    let quote = extract_quote(&text);
     let quote_pos = text.find(quote).unwrap();
     let prefix = text[..quote_pos].to_lowercase();
     let body = &text[(quote_pos + quote.len())..(text.len() - quote.len())];

src/pydocstyle/plugins.rs

@@ -210,7 +210,7 @@ pub fn blank_before_after_function(checker: &mut Checker, definition: &Definitio
                         .count();
 
                     // Avoid D202 violations for blank lines followed by inner functions or classes.
-                    if blank_lines_after == 1 && INNER_FUNCTION_OR_CLASS_REGEX.is_match(after) {
+                    if blank_lines_after == 1 && INNER_FUNCTION_OR_CLASS_REGEX.is_match(&after) {
                         return;
                     }
 
@@ -387,7 +387,7 @@ pub fn indent(checker: &mut Checker, definition: &Definition) {
                 return;
             }
 
-            let docstring_indent = helpers::indentation(checker, docstring).to_string();
+            let docstring_indent = helpers::indentation(checker, docstring);
             let mut has_seen_tab = docstring_indent.contains('\t');
             let mut is_over_indented = true;
             let mut over_indented_lines = vec![];
@@ -537,7 +537,7 @@ pub fn newline_after_last_paragraph(checker: &mut Checker, definition: &Definiti
                                 // Insert a newline just before the end-quote(s).
                                 let content = format!(
                                     "\n{}",
-                                    helpers::clean(helpers::indentation(checker, docstring))
+                                    helpers::clean(&helpers::indentation(checker, docstring))
                                 );
                                 check.amend(Fix::insertion(
                                     content,
@@ -916,7 +916,7 @@ fn blanks_and_section_underline(
                 // Add a dashed line (of the appropriate length) under the section header.
                 let content = format!(
                     "{}{}\n",
-                    helpers::clean(helpers::indentation(checker, docstring)),
+                    helpers::clean(&helpers::indentation(checker, docstring)),
                     "-".repeat(context.section_name.len())
                 );
                 check.amend(Fix::insertion(
@@ -950,7 +950,7 @@ fn blanks_and_section_underline(
                 // Add a dashed line (of the appropriate length) under the section header.
                 let content = format!(
                     "{}{}\n",
-                    helpers::clean(helpers::indentation(checker, docstring)),
+                    helpers::clean(&helpers::indentation(checker, docstring)),
                     "-".repeat(context.section_name.len())
                 );
                 check.amend(Fix::insertion(
@@ -1026,7 +1026,7 @@ fn blanks_and_section_underline(
                     // Replace the existing underline with a line of the appropriate length.
                     let content = format!(
                         "{}{}\n",
-                        helpers::clean(helpers::indentation(checker, docstring)),
+                        helpers::clean(&helpers::indentation(checker, docstring)),
                         "-".repeat(context.section_name.len())
                     );
                     check.amend(Fix::replacement(
@@ -1054,7 +1054,7 @@ fn blanks_and_section_underline(
 
         if checker.settings.enabled.contains(&CheckCode::D215) {
             let leading_space = helpers::leading_space(non_empty_line);
-            let indentation = helpers::indentation(checker, docstring).to_string();
+            let indentation = helpers::indentation(checker, docstring);
             if leading_space.len() > indentation.len() {
                 let mut check = Check::new(
                     CheckKind::SectionUnderlineNotOverIndented(context.section_name.to_string()),
@@ -1195,7 +1195,7 @@ fn common_section(
 
     if checker.settings.enabled.contains(&CheckCode::D214) {
         let leading_space = helpers::leading_space(context.line);
-        let indentation = helpers::indentation(checker, docstring).to_string();
+        let indentation = helpers::indentation(checker, docstring);
         if leading_space.len() > indentation.len() {
             let mut check = Check::new(
                 CheckKind::SectionNotOverIndented(context.section_name.to_string()),