Fix behavior around comments embedded in imports

README.md

@@ -34,6 +34,7 @@ Since then more critical features & fixes have been added, and the options have
     - [`importOrderTypeScriptVersion`](#importordertypescriptversion)
     - [`importOrderParserPlugins`](#importorderparserplugins)
   - [Prevent imports from being sorted](#prevent-imports-from-being-sorted)
+  - [Comments](#comments)
 - [FAQ / Troubleshooting](#faq--troubleshooting)
 - [Compatibility](#compatibility)
 - [Contribution](#contribution)
@@ -271,6 +272,15 @@ This will keep the `zealand` import at the top instead of moving it below the `a
 entire import statements can be ignored, line comments (`// prettier-ignore`) are recommended over inline comments
 (`/* prettier-ignore */`).
 
+### Comments
+
+We make the following attempts at keeping comments in your imports clean:
+
+- If you have one or more comments at the top of the file, we will keep them at the top as long as there is a blank line before your first import statement.
+- Comments on lines after the final import statement will not be moved. (Runtime-code between imports will be moved below all the imports).
+- In general, if you place a comment on the same line as an Import `Declaration` or `*Specifier`, we will keep it attached to that same specifier if it moves around.
+- Other comments are preserved, and are generally considered "leading" comments for the subsequent Import `Declaration` or `*Specifier`.
+
 ## FAQ / Troubleshooting
 
 Having some trouble or an issue? You can check [FAQ / Troubleshooting section](./docs/TROUBLESHOOTING.md).

package.json

@@ -59,14 +59,12 @@
         "@babel/traverse": "^7.21.5",
         "@babel/types": "^7.21.5",
         "lodash.clone": "^4.5.0",
-        "lodash.isequal": "^4.5.0",
         "semver": "^7.5.0"
     },
     "devDependencies": {
         "@types/babel__generator": "^7.6.4",
         "@types/babel__traverse": "^7.18.3",
         "@types/lodash.clone": "4.5.7",
-        "@types/lodash.isequal": "4.5.6",
         "@types/node": "^18.15.13",
         "@types/prettier": "^2.7.2",
         "@types/semver": "^7.3.13",

src/constants.ts

@@ -34,6 +34,18 @@ export const TYPES_SPECIAL_WORD = '<TYPES>';
 const PRETTIER_PLUGIN_SORT_IMPORTS_NEW_LINE =
     'PRETTIER_PLUGIN_SORT_IMPORTS_NEW_LINE';
 
+/** Use this to force a newline at top-level scope (good for newlines generated between import blocks) */
 export const newLineNode = expressionStatement(
     stringLiteral(PRETTIER_PLUGIN_SORT_IMPORTS_NEW_LINE),
 );
+/** Use this if you want to force a newline, but you're attaching to leading/inner/trailing Comments */
+export const forceANewlineUsingACommentStatement = () => ({
+    type: 'CommentLine' as const,
+    value: 'PRETTIER_PLUGIN_SORT_IMPORTS_NEWLINE_COMMENT',
+    start: -1,
+    end: -1,
+    loc: { start: { line: -1, column: -1 }, end: { line: -1, column: -1 } },
+});
+
+export const injectNewlinesRegex =
+    /("PRETTIER_PLUGIN_SORT_IMPORTS_NEW_LINE";|\/\/PRETTIER_PLUGIN_SORT_IMPORTS_NEWLINE_COMMENT)/gi;

src/preprocessors/preprocessor.ts

@@ -31,6 +31,7 @@ export function preprocessor(code: string, options: PrettierOptions): string {
     const allOriginalImportNodes: ImportDeclaration[] = [];
     const parserOptions: ParserOptions = {
         sourceType: 'module',
+        attachComment: true,
         plugins: getExperimentalParserPlugins(importOrderParserPlugins),
     };
 

src/types.ts

@@ -1,4 +1,11 @@
-import { ExpressionStatement, ImportDeclaration } from '@babel/types';
+import {
+    type EmptyStatement,
+    type ExpressionStatement,
+    type ImportDeclaration,
+    type ImportDefaultSpecifier,
+    type ImportNamespaceSpecifier,
+    type ImportSpecifier,
+} from '@babel/types';
 import { RequiredOptions } from 'prettier';
 
 import { PluginConfig } from '../types';
@@ -28,7 +35,16 @@ export interface ImportChunk {
 }
 
 export type ImportGroups = Record<string, ImportDeclaration[]>;
-export type ImportOrLine = ImportDeclaration | ExpressionStatement;
+export type ImportOrLine =
+    | ImportDeclaration
+    | ExpressionStatement
+    | EmptyStatement;
+
+export type SomeSpecifier =
+    | ImportSpecifier
+    | ImportDefaultSpecifier
+    | ImportNamespaceSpecifier;
+export type ImportRelated = ImportOrLine | SomeSpecifier;
 
 export type GetSortedNodes = (
     nodes: ImportDeclaration[],

src/utils/__tests__/adjust-comments-on-sorted-nodes.spec.ts

@@ -23,13 +23,18 @@ test('it preserves the single leading comment for each import declaration', () =
     expect(importNodes).toHaveLength(3);
     const finalNodes = [importNodes[2], importNodes[1], importNodes[0]];
     const adjustedNodes = adjustCommentsOnSortedNodes(importNodes, finalNodes);
-    expect(adjustedNodes).toHaveLength(3);
-    expect(leadingComments(adjustedNodes[0])).toEqual([' comment a']);
+    expect(adjustedNodes).toHaveLength(4);
+    // First node is a dummy EmptyStatement
+    expect(adjustedNodes[0].type).toEqual('EmptyStatement');
+    expect(leadingComments(adjustedNodes[0])).toEqual([]);
     expect(trailingComments(adjustedNodes[0])).toEqual([]);
-    expect(leadingComments(adjustedNodes[1])).toEqual([' comment b']);
+    expect(leadingComments(adjustedNodes[1])).toEqual([' comment a']);
     expect(trailingComments(adjustedNodes[1])).toEqual([]);
-    expect(leadingComments(adjustedNodes[2])).toEqual([]);
+    expect(leadingComments(adjustedNodes[2])).toEqual([' comment b']);
     expect(trailingComments(adjustedNodes[2])).toEqual([]);
+    // Import from "c" has no leading comment, and the trailing was kept with "b"
+    expect(leadingComments(adjustedNodes[3])).toEqual([]);
+    expect(trailingComments(adjustedNodes[3])).toEqual([]);
 });
 
 test('it preserves multiple leading comments for each import declaration', () => {
@@ -47,24 +52,24 @@ test('it preserves multiple leading comments for each import declaration', () =>
     expect(importNodes).toHaveLength(3);
     const finalNodes = [importNodes[2], importNodes[1], importNodes[0]];
     const adjustedNodes = adjustCommentsOnSortedNodes(importNodes, finalNodes);
-    expect(adjustedNodes).toHaveLength(3);
-    expect(leadingComments(adjustedNodes[0])).toEqual([
+    expect(adjustedNodes).toHaveLength(4);
+    expect(leadingComments(adjustedNodes[1])).toEqual([
         ' comment a1',
         ' comment a2',
         ' comment a3',
     ]);
-    expect(trailingComments(adjustedNodes[0])).toEqual([]);
-    expect(leadingComments(adjustedNodes[1])).toEqual([
+    expect(trailingComments(adjustedNodes[1])).toEqual([]);
+    expect(leadingComments(adjustedNodes[2])).toEqual([
         ' comment b1',
         ' comment b2',
         ' comment b3',
     ]);
-    expect(trailingComments(adjustedNodes[1])).toEqual([]);
-    expect(leadingComments(adjustedNodes[2])).toEqual([]);
     expect(trailingComments(adjustedNodes[2])).toEqual([]);
+    expect(leadingComments(adjustedNodes[3])).toEqual([]);
+    expect(trailingComments(adjustedNodes[3])).toEqual([]);
 });
 
-test('it does not move comments at before all import declarations', () => {
+test('it does not move comments more than one line before all import declarations', () => {
     const importNodes = getImportNodes(`
     // comment c1
     // comment c2
@@ -75,16 +80,16 @@ test('it does not move comments at before all import declarations', () => {
     expect(importNodes).toHaveLength(3);
     const finalNodes = [importNodes[2], importNodes[1], importNodes[0]];
     const adjustedNodes = adjustCommentsOnSortedNodes(importNodes, finalNodes);
-    expect(adjustedNodes).toHaveLength(3);
-    expect(leadingComments(adjustedNodes[0])).toEqual([
-        ' comment c1',
-        ' comment c2',
-    ]);
-    expect(trailingComments(adjustedNodes[0])).toEqual([]);
-    expect(leadingComments(adjustedNodes[1])).toEqual([]);
-    expect(trailingComments(adjustedNodes[1])).toEqual([]);
+    expect(adjustedNodes).toHaveLength(4);
+    // Comment c1 is more than one line above the first import, so it stays with the top-of-file
+    expect(leadingComments(adjustedNodes[0])).toEqual([' comment c1']);
+
     expect(leadingComments(adjustedNodes[2])).toEqual([]);
     expect(trailingComments(adjustedNodes[2])).toEqual([]);
+    expect(trailingComments(adjustedNodes[3])).toEqual([]);
+
+    // Comment c2 is attached to import from "c"
+    expect(leadingComments(adjustedNodes[3])).toEqual([' comment c2']);
 });
 
 test('it does not affect comments after all import declarations', () => {
@@ -98,11 +103,11 @@ test('it does not affect comments after all import declarations', () => {
     expect(importNodes).toHaveLength(3);
     const finalNodes = [importNodes[2], importNodes[1], importNodes[0]];
     const adjustedNodes = adjustCommentsOnSortedNodes(importNodes, finalNodes);
-    expect(adjustedNodes).toHaveLength(3);
-    expect(leadingComments(adjustedNodes[0])).toEqual([]);
-    expect(trailingComments(adjustedNodes[0])).toEqual([]);
+    expect(adjustedNodes).toHaveLength(4);
     expect(leadingComments(adjustedNodes[1])).toEqual([]);
     expect(trailingComments(adjustedNodes[1])).toEqual([]);
     expect(leadingComments(adjustedNodes[2])).toEqual([]);
     expect(trailingComments(adjustedNodes[2])).toEqual([]);
+    expect(leadingComments(adjustedNodes[3])).toEqual([]);
+    expect(trailingComments(adjustedNodes[3])).toEqual([]);
 });

src/utils/__tests__/get-code-from-ast.spec.ts

@@ -6,9 +6,7 @@ import { getImportNodes } from '../get-import-nodes';
 import { getSortedNodes } from '../get-sorted-nodes';
 
 test('sorts imports correctly', () => {
-    const code = `// first comment
-// second comment
-import z from 'z';
+    const code = `import z from 'z';
 import c from 'c';
 import g from 'g';
 import t from 't';
@@ -26,10 +24,7 @@ import a from 'a';
         directives: [],
     });
     expect(format(formatted, { parser: 'babel' })).toEqual(
-        `// first comment
-// second comment
-
-import a from "a";
+        `import a from "a";
 import c from "c";
 import g from "g";
 import k from "k";
@@ -40,9 +35,7 @@ import z from "z";
 });
 
 test('merges duplicate imports correctly', () => {
-    const code = `// first comment
-// second comment
-import z from 'z';
+    const code = `import z from 'z';
 import c from 'c';
 import g from 'g';
 import t from 't';
@@ -64,10 +57,7 @@ import type {See} from 'c';
         directives: [],
     });
     expect(format(formatted, { parser: 'babel' })).toEqual(
-        `// first comment
-// second comment
-
-import a, { b, type Bee } from "a";
+        `import a, { b, type Bee } from "a";
 import c, { type C, type See } from "c";
 import g from "g";
 import k from "k";

src/utils/__tests__/get-sorted-nodes.spec.ts

@@ -6,8 +6,7 @@ import { getSortedNodes } from '../get-sorted-nodes';
 import { getSortedNodesModulesNames } from '../get-sorted-nodes-modules-names';
 import { getSortedNodesNamesAndNewlines } from '../get-sorted-nodes-names-and-newlines';
 
-const code = `// first comment
-// second comment
+const code = `
 import "se3";
 import z from 'z';
 import c, { cD } from 'c';
@@ -51,13 +50,14 @@ test('it returns all sorted nodes, preserving the order side effect nodes', () =
         'se2',
         '',
     ]);
-    expect(
-        sorted
-            .filter((node) => node.type === 'ImportDeclaration')
-            .map((importDeclaration) =>
-                getSortedNodesModulesNames(importDeclaration.specifiers),
-            ),
-    ).toEqual([
+
+    const result2 = sorted
+        .filter((node) => node.type === 'ImportDeclaration')
+        .map((importDeclaration) =>
+            getSortedNodesModulesNames(importDeclaration.specifiers),
+        );
+
+    expect(result2).toEqual([
         [],
         ['c', 'cD'],
         ['g'],

src/utils/adjust-comments-on-sorted-nodes.ts

@@ -1,43 +1,50 @@
-import {
-    addComments,
-    removeComments,
-    type ImportDeclaration,
-} from '@babel/types';
-import clone from 'lodash.clone';
-import isEqual from 'lodash.isequal';
+import { removeComments, type ImportDeclaration } from '@babel/types';
 
 import { ImportOrLine } from '../types';
+import {
+    attachCommentsToOutputNodes,
+    getCommentRegistryFromImportDeclarations,
+} from './get-comment-registry';
 
 /**
  * Takes the original nodes before sorting and the final nodes after sorting.
  * Adjusts the comments on the final nodes so that they match the comments as
  * they were in the original nodes.
- * @param nodes A list of nodes in the order as they were originally.
+ * @param originalDeclarationNodes A list of nodes in the order as they were originally.
  * @param finalNodes The same set of nodes, but in the final sorting order.
  * @returns A copied and adjusted set of nodes, containing comments
  */
 export const adjustCommentsOnSortedNodes = (
-    nodes: ImportDeclaration[],
+    originalDeclarationNodes: ImportDeclaration[],
     finalNodes: ImportOrLine[],
 ) => {
-    // We will mutate a copy of the finalNodes, and extract comments from the original
-    const finalNodesClone = finalNodes.map(clone);
-
-    const firstNodesComments = nodes[0].leadingComments;
-
-    // Remove all comments from sorted nodes
-    finalNodesClone.forEach(removeComments);
+    const outputNodes: ImportDeclaration[] = finalNodes.filter(
+        (n) => n.type === 'ImportDeclaration',
+    ) as ImportDeclaration[];
+    if (originalDeclarationNodes.length === 0 || outputNodes.length === 0) {
+        // Nothing to do, because there are no ImportDeclarations!
+        return finalNodes;
+    }
 
-    // insert comments other than the first comments
-    finalNodesClone.forEach((node, index) => {
-        if (isEqual(nodes[0].loc, node.loc)) return;
+    const registry = getCommentRegistryFromImportDeclarations({
+        outputNodes,
+        firstImport: originalDeclarationNodes[0],
+    });
 
-        addComments(node, 'leading', finalNodes[index].leadingComments || []);
+    // Make a copy of the nodes for easier debugging & remove the existing comments to reattach them
+    // (removeComments clones the nodes internally, so we don't need to do that ourselves)
+    const finalNodesClone = finalNodes.map((n) => {
+        const noDirectCommentsNode = removeComments(n);
+        if (noDirectCommentsNode.type === 'ImportDeclaration') {
+            // Remove comments isn't recursive, so we need to clone/modify the specifiers manually
+            noDirectCommentsNode.specifiers = (
+                noDirectCommentsNode.specifiers || []
+            ).map((s) => removeComments(s));
+        }
+        return noDirectCommentsNode;
     });
 
-    if (firstNodesComments) {
-        addComments(finalNodesClone[0], 'leading', firstNodesComments);
-    }
+    attachCommentsToOutputNodes(registry, finalNodesClone);
 
     return finalNodesClone;
 };

src/utils/get-all-comments-from-nodes.ts

@@ -5,13 +5,32 @@ import type {
     Statement,
 } from '@babel/types';
 
-export const getAllCommentsFromNodes = (nodes: (Directive | Statement)[]) =>
+import { SomeSpecifier } from '../types';
+
+export const getAllCommentsFromNodes = (
+    nodes: (Directive | Statement | SomeSpecifier)[],
+) =>
     nodes.reduce((acc, node) => {
         if (
             Array.isArray(node.leadingComments) &&
             node.leadingComments.length > 0
         ) {
             acc = [...acc, ...node.leadingComments];
         }
+        if (
+            Array.isArray(node.innerComments) &&
+            node.innerComments.length > 0
+        ) {
+            acc = [...acc, ...node.innerComments];
+        }
+        if (
+            Array.isArray(node.trailingComments) &&
+            node.trailingComments.length > 0
+        ) {
+            acc = [...acc, ...node.trailingComments];
+        }
+        if (node.type === 'ImportDeclaration') {
+            acc = [...acc, ...getAllCommentsFromNodes(node.specifiers)];
+        }
         return acc;
     }, [] as (CommentBlock | CommentLine)[]);