Coming together

foo.py → resources/test/fixtures/isort/comments.py

@@ -4,6 +4,7 @@
 
 # Comment 3a
 import C
+
 # Comment 3b
 import C
 
@@ -13,12 +14,12 @@
 
 # Comment 6
 from A import (
-	a,  # Comment 7
-	b,
-	c,  # Comment 8
+    a,  # Comment 7
+    b,
+    c,  # Comment 8
 )
 from A import (
-	a,  # Comment 9
-	b,  # Comment 10
-	c,  # Comment 11
+    a,  # Comment 9
+    b,  # Comment 10
+    c,  # Comment 11
 )

resources/test/fixtures/isort/type_comments.py

@@ -0,0 +1,2 @@
+import A  # type: ignore
+from B import C  # type: ignore

src/isort/comments.rs

@@ -17,12 +17,10 @@ pub struct Comment<'a> {
 
 /// Collect all comments in an import block.
 pub fn collect_comments<'a>(range: &Range, locator: &'a SourceCodeLocator) -> Vec<Comment<'a>> {
-    let contents = locator.slice_source_code_range(&range);
-    println!("conents = '{}'", contents);
+    let contents = locator.slice_source_code_range(range);
     lexer::make_tokenizer(&contents)
         .flatten()
         .filter_map(|(start, tok, end)| {
-            println!("{:?} {:?} {:?}", start, tok, end);
             if matches!(tok, Tok::Comment) {
                 let start = helpers::to_absolute(&start, &range.location);
                 let end = helpers::to_absolute(&end, &range.location);
@@ -41,5 +39,5 @@ pub fn collect_comments<'a>(range: &Range, locator: &'a SourceCodeLocator) -> Ve
         .collect()
 }
 
-// We have to assign each comment to a node. Then when we collect, we merge nodes, or something.
-// Comments can either be above, or to the right of a node.
+// We have to assign each comment to a node. Then when we collect, we merge
+// nodes, or something. Comments can either be above, or to the right of a node.

src/isort/format.rs

@@ -0,0 +1,150 @@
+use ropey::RopeBuilder;
+
+use crate::isort::types::{AliasData, CommentSet, ImportFromData, Importable};
+
+// Hard-code four-space indentation for the imports themselves, to match Black.
+const INDENT: &str = "    ";
+
+/// Add a plain import statement to the `RopeBuilder`.
+pub fn format_import(output: &mut RopeBuilder, alias: &AliasData, comments: &CommentSet) {
+    for comment in &comments.atop {
+        output.append(&format!("{}\n", comment));
+    }
+    if let Some(asname) = alias.asname {
+        output.append(&format!("import {} as {}", alias.name, asname));
+    } else {
+        output.append(&format!("import {}", alias.name));
+    }
+    for comment in &comments.inline {
+        output.append(&format!("  {}", comment));
+    }
+    output.append("\n");
+}
+
+/// Add an import-from statement to the `RopeBuilder`.
+pub fn format_import_from(
+    output: &mut RopeBuilder,
+    import_from: &ImportFromData,
+    comments: &CommentSet,
+    aliases: &[(AliasData, CommentSet)],
+    line_length: &usize,
+) {
+    // We can only inline if: (1) none of the aliases have atop comments, and (3)
+    // only the last alias (if any) has inline comments.
+    if aliases
+        .iter()
+        .all(|(_, CommentSet { atop, .. })| atop.is_empty())
+        && aliases
+            .iter()
+            .rev()
+            .skip(1)
+            .all(|(_, CommentSet { inline, .. })| inline.is_empty())
+    {
+        let single_line = format_single_line(import_from, comments, aliases);
+        // If the import fits on a single line (excluding the newline character at the
+        // end, which doesn't count towards the line length), return it.
+        if single_line.len() <= *line_length + 1 {
+            output.append(&single_line);
+            return;
+        }
+    }
+
+    output.append(&format_multi_line(import_from, comments, aliases));
+}
+
+/// Format an import-from statement in single-line format.
+///
+/// This method assumes that the output source code is syntactically valid.
+fn format_single_line(
+    import_from: &ImportFromData,
+    comments: &CommentSet,
+    aliases: &[(AliasData, CommentSet)],
+) -> String {
+    let mut output = String::new();
+
+    for comment in &comments.atop {
+        output.push_str(comment);
+        output.push('\n');
+    }
+
+    output.push_str(&format!("from {} import ", import_from.module_name()));
+
+    for (index, (AliasData { name, asname }, comments)) in aliases.iter().enumerate() {
+        for comment in &comments.atop {
+            output.push_str(comment);
+            output.push('\n');
+        }
+        if let Some(asname) = asname {
+            output.push_str(name);
+            output.push_str(" as ");
+            output.push_str(asname);
+        } else {
+            output.push_str(name);
+        }
+        if index < aliases.len() - 1 {
+            output.push_str(", ");
+        }
+
+        for comment in &comments.inline {
+            output.push_str(&format!("  {}", comment));
+        }
+    }
+
+    for comment in &comments.inline {
+        output.push_str(&format!("  {}", comment));
+    }
+
+    output.push('\n');
+
+    output
+}
+
+/// Format an import-from statement in multi-line format.
+fn format_multi_line(
+    import_from: &ImportFromData,
+    comments: &CommentSet,
+    aliases: &[(AliasData, CommentSet)],
+) -> String {
+    let mut output = String::new();
+
+    for comment in &comments.atop {
+        output.push_str(comment);
+        output.push('\n');
+    }
+
+    output.push_str(&format!("from {} import ", import_from.module_name()));
+    output.push('(');
+    for comment in &comments.inline {
+        output.push_str(&format!("  {}", comment));
+    }
+    output.push('\n');
+
+    for (AliasData { name, asname }, comments) in aliases {
+        for comment in &comments.atop {
+            output.push_str(INDENT);
+            output.push_str(comment);
+            output.push('\n');
+        }
+        output.push_str(INDENT);
+        if let Some(asname) = asname {
+            output.push_str(name);
+            output.push_str(" as ");
+            output.push_str(asname);
+        } else {
+            output.push_str(name);
+        }
+        output.push(',');
+
+        for comment in &comments.inline {
+            output.push(' ');
+            output.push(' ');
+            output.push_str(comment);
+        }
+        output.push('\n');
+    }
+
+    output.push(')');
+    output.push('\n');
+
+    output
+}