fix 57215 -- add support for import attributes to OrganizeImports (#57250)

iisaduan · web-flow · commit ee2090d7f2af · 2024-02-01T10:32:58.000-08:00
diff --git a/src/services/organizeImports.ts b/src/services/organizeImports.ts
@@ -59,6 +59,7 @@ import {
     Scanner,
     setEmitFlags,
     some,
+    sort,
     SortKind,
     SourceFile,
     stableSort,
@@ -322,96 +323,111 @@ function coalesceImportsWorker(importGroup: readonly ImportDeclaration[], compar
         return importGroup;
     }
 
-    const { importWithoutClause, typeOnlyImports, regularImports } = getCategorizedImports(importGroup);
-    const coalescedImports: ImportDeclaration[] = [];
-
-    if (importWithoutClause) {
-        coalescedImports.push(importWithoutClause);
-    }
-
-    for (const group of [regularImports, typeOnlyImports]) {
-        const isTypeOnly = group === typeOnlyImports;
-        const { defaultImports, namespaceImports, namedImports } = group;
-        // Normally, we don't combine default and namespace imports, but it would be silly to
-        // produce two import declarations in this special case.
-        if (!isTypeOnly && defaultImports.length === 1 && namespaceImports.length === 1 && namedImports.length === 0) {
-            // Add the namespace import to the existing default ImportDeclaration.
-            const defaultImport = defaultImports[0];
-            coalescedImports.push(
-                updateImportDeclarationAndClause(defaultImport, defaultImport.importClause.name, namespaceImports[0].importClause.namedBindings),
-            );
-
-            continue;
+    const importGroupsByAttributes = groupBy(importGroup, decl => {
+        if (decl.attributes) {
+            let attrs = decl.attributes.token + " ";
+            for (const x of sort(decl.attributes.elements, (x, y) => compareStringsCaseSensitive(x.name.text, y.name.text))) {
+                attrs += x.name.text + ":";
+                attrs += isStringLiteralLike(x.value) ? `"${x.value.text}"` : x.value.getText() + " ";
+            }
+            return attrs;
         }
+        return "";
+    });
 
-        const sortedNamespaceImports = stableSort(namespaceImports, (i1, i2) => comparer(i1.importClause.namedBindings.name.text, i2.importClause.namedBindings.name.text));
+    const coalescedImports: ImportDeclaration[] = [];
+    for (const attribute in importGroupsByAttributes) {
+        const importGroupSameAttrs = importGroupsByAttributes[attribute] as ImportDeclaration[];
+        const { importWithoutClause, typeOnlyImports, regularImports } = getCategorizedImports(importGroupSameAttrs);
+
+        if (importWithoutClause) {
+            coalescedImports.push(importWithoutClause);
+        }
+
+        for (const group of [regularImports, typeOnlyImports]) {
+            const isTypeOnly = group === typeOnlyImports;
+            const { defaultImports, namespaceImports, namedImports } = group;
+            // Normally, we don't combine default and namespace imports, but it would be silly to
+            // produce two import declarations in this special case.
+            if (!isTypeOnly && defaultImports.length === 1 && namespaceImports.length === 1 && namedImports.length === 0) {
+                // Add the namespace import to the existing default ImportDeclaration.
+                const defaultImport = defaultImports[0];
+                coalescedImports.push(
+                    updateImportDeclarationAndClause(defaultImport, defaultImport.importClause.name, namespaceImports[0].importClause.namedBindings),
+                );
 
-        for (const namespaceImport of sortedNamespaceImports) {
-            // Drop the name, if any
-            coalescedImports.push(
-                updateImportDeclarationAndClause(namespaceImport, /*name*/ undefined, namespaceImport.importClause.namedBindings),
-            );
-        }
+                continue;
+            }
 
-        const firstDefaultImport = firstOrUndefined(defaultImports);
-        const firstNamedImport = firstOrUndefined(namedImports);
-        const importDecl = firstDefaultImport ?? firstNamedImport;
-        if (!importDecl) {
-            continue;
-        }
+            const sortedNamespaceImports = stableSort(namespaceImports, (i1, i2) => comparer(i1.importClause.namedBindings.name.text, i2.importClause.namedBindings.name.text));
 
-        let newDefaultImport: Identifier | undefined;
-        const newImportSpecifiers: ImportSpecifier[] = [];
-        if (defaultImports.length === 1) {
-            newDefaultImport = defaultImports[0].importClause.name;
-        }
-        else {
-            for (const defaultImport of defaultImports) {
-                newImportSpecifiers.push(
-                    factory.createImportSpecifier(/*isTypeOnly*/ false, factory.createIdentifier("default"), defaultImport.importClause.name),
+            for (const namespaceImport of sortedNamespaceImports) {
+                // Drop the name, if any
+                coalescedImports.push(
+                    updateImportDeclarationAndClause(namespaceImport, /*name*/ undefined, namespaceImport.importClause.namedBindings),
                 );
             }
-        }
 
-        newImportSpecifiers.push(...getNewImportSpecifiers(namedImports));
+            const firstDefaultImport = firstOrUndefined(defaultImports);
+            const firstNamedImport = firstOrUndefined(namedImports);
+            const importDecl = firstDefaultImport ?? firstNamedImport;
+            if (!importDecl) {
+                continue;
+            }
 
-        const sortedImportSpecifiers = factory.createNodeArray(
-            sortSpecifiers(newImportSpecifiers, comparer, preferences),
-            firstNamedImport?.importClause.namedBindings.elements.hasTrailingComma,
-        );
+            let newDefaultImport: Identifier | undefined;
+            const newImportSpecifiers: ImportSpecifier[] = [];
+            if (defaultImports.length === 1) {
+                newDefaultImport = defaultImports[0].importClause.name;
+            }
+            else {
+                for (const defaultImport of defaultImports) {
+                    newImportSpecifiers.push(
+                        factory.createImportSpecifier(/*isTypeOnly*/ false, factory.createIdentifier("default"), defaultImport.importClause.name),
+                    );
+                }
+            }
 
-        const newNamedImports = sortedImportSpecifiers.length === 0
-            ? newDefaultImport
-                ? undefined
-                : factory.createNamedImports(emptyArray)
-            : firstNamedImport
-            ? factory.updateNamedImports(firstNamedImport.importClause.namedBindings, sortedImportSpecifiers)
-            : factory.createNamedImports(sortedImportSpecifiers);
-
-        if (
-            sourceFile &&
-            newNamedImports &&
-            firstNamedImport?.importClause.namedBindings &&
-            !rangeIsOnSingleLine(firstNamedImport.importClause.namedBindings, sourceFile)
-        ) {
-            setEmitFlags(newNamedImports, EmitFlags.MultiLine);
-        }
+            newImportSpecifiers.push(...getNewImportSpecifiers(namedImports));
 
-        // Type-only imports are not allowed to mix default, namespace, and named imports in any combination.
-        // We could rewrite a default import as a named import (`import { default as name }`), but we currently
-        // choose not to as a stylistic preference.
-        if (isTypeOnly && newDefaultImport && newNamedImports) {
-            coalescedImports.push(
-                updateImportDeclarationAndClause(importDecl, newDefaultImport, /*namedBindings*/ undefined),
-            );
-            coalescedImports.push(
-                updateImportDeclarationAndClause(firstNamedImport ?? importDecl, /*name*/ undefined, newNamedImports),
-            );
-        }
-        else {
-            coalescedImports.push(
-                updateImportDeclarationAndClause(importDecl, newDefaultImport, newNamedImports),
+            const sortedImportSpecifiers = factory.createNodeArray(
+                sortSpecifiers(newImportSpecifiers, comparer, preferences),
+                firstNamedImport?.importClause.namedBindings.elements.hasTrailingComma,
             );
+
+            const newNamedImports = sortedImportSpecifiers.length === 0
+                ? newDefaultImport
+                    ? undefined
+                    : factory.createNamedImports(emptyArray)
+                : firstNamedImport
+                ? factory.updateNamedImports(firstNamedImport.importClause.namedBindings, sortedImportSpecifiers)
+                : factory.createNamedImports(sortedImportSpecifiers);
+
+            if (
+                sourceFile &&
+                newNamedImports &&
+                firstNamedImport?.importClause.namedBindings &&
+                !rangeIsOnSingleLine(firstNamedImport.importClause.namedBindings, sourceFile)
+            ) {
+                setEmitFlags(newNamedImports, EmitFlags.MultiLine);
+            }
+
+            // Type-only imports are not allowed to mix default, namespace, and named imports in any combination.
+            // We could rewrite a default import as a named import (`import { default as name }`), but we currently
+            // choose not to as a stylistic preference.
+            if (isTypeOnly && newDefaultImport && newNamedImports) {
+                coalescedImports.push(
+                    updateImportDeclarationAndClause(importDecl, newDefaultImport, /*namedBindings*/ undefined),
+                );
+                coalescedImports.push(
+                    updateImportDeclarationAndClause(firstNamedImport ?? importDecl, /*name*/ undefined, newNamedImports),
+                );
+            }
+            else {
+                coalescedImports.push(
+                    updateImportDeclarationAndClause(importDecl, newDefaultImport, newNamedImports),
+                );
+            }
         }
     }
 
diff --git a/tests/cases/fourslash/organizeImportsAttributes.ts b/tests/cases/fourslash/organizeImportsAttributes.ts
@@ -0,0 +1,20 @@
+/// <reference path="fourslash.ts" />
+
+//// import { A } from "./file";
+//// import { type B } from "./file";
+//// import { C } from "./file" assert { type: "a" };
+//// import { A as D } from "./file" assert { type: "b" };
+//// import { E } from "./file" with { type: "a" };
+//// import { A as F } from "./file" with { type: "b" };
+//// 
+//// type G = A | B | C | D | E | F;
+
+verify.organizeImports(
+`import { A, type B } from "./file";
+import { C } from "./file" assert { type: "a" };
+import { A as D } from "./file" assert { type: "b" };
+import { E } from "./file" with { type: "a" };
+import { A as F } from "./file" with { type: "b" };
+
+type G = A | B | C | D | E | F;`);
+    
diff --git a/tests/cases/fourslash/organizeImportsAttributes2.ts b/tests/cases/fourslash/organizeImportsAttributes2.ts
@@ -0,0 +1,20 @@
+/// <reference path="fourslash.ts" />
+
+////import { A } from "./a";
+////import { C } from "./a" assert { type: "a" };
+////import { Z } from "./z";
+////import { A as D } from "./a" assert { type: "b" };
+////import { E } from "./a" with { type: "a" };
+////import { F } from "./a" assert { type: "a" };
+////import { B } from "./a";
+////
+////export type G = A | B | C | D | E | F | Z;
+
+verify.organizeImports(
+`import { A, B } from "./a";
+import { C, F } from "./a" assert { type: "a" };
+import { A as D } from "./a" assert { type: "b" };
+import { E } from "./a" with { type: "a" };
+import { Z } from "./z";
+
+export type G = A | B | C | D | E | F | Z;`);
diff --git a/tests/cases/fourslash/organizeImportsAttributes3.ts b/tests/cases/fourslash/organizeImportsAttributes3.ts
@@ -0,0 +1,20 @@
+/// <reference path="fourslash.ts" />
+
+////import { A } from "./a";
+////import { C } from "./a" assert {      type: "a" };
+////import { Z } from "./z";
+////import { A as D } from "./a" assert    { type: "b" };
+////import { E } from "./a" assert { type: /* comment*/ "a"              };
+////import { F } from "./a" assert     {type: "a" };
+////import { Y } from "./a"   assert{ type: "b" /* comment*/};
+////import { B } from "./a";
+////
+////export type G = A | B | C | D | E | F | Y | Z;
+
+verify.organizeImports(
+`import { A, B } from "./a";
+import { C, E, F } from "./a" assert { type: "a" };
+import { A as D, Y } from "./a" assert { type: "b" };
+import { Z } from "./z";
+
+export type G = A | B | C | D | E | F | Y | Z;`);
diff --git a/tests/cases/fourslash/organizeImportsAttributes4.ts b/tests/cases/fourslash/organizeImportsAttributes4.ts
@@ -0,0 +1,21 @@
+/// <reference path="fourslash.ts" />
+
+////import { A } from "./a" assert { foo: "foo", bar: "bar" };
+////import { B } from "./a" assert { bar: "bar", foo: "foo" };
+////import { D } from "./a" assert { bar: "foo", foo: "bar" };
+////import { E } from "./a" assert { foo: 'bar', bar: "foo" };
+////import { C } from "./a" assert { foo: "bar", bar: "foo" };
+////import { F } from "./a" assert { foo: "42" };
+////import { Y } from "./a" assert { foo: 42 };
+////import { Z } from "./a" assert { foo: "42" };
+////
+////export type G = A | B | C | D | E | F | Y | Z;
+
+
+verify.organizeImports(
+`import { A, B } from "./a" assert { foo: "foo", bar: "bar" };
+import { C, D, E } from "./a" assert { bar: "foo", foo: "bar" };
+import { F, Z } from "./a" assert { foo: "42" };
+import { Y } from "./a" assert { foo: 42 };
+
+export type G = A | B | C | D | E | F | Y | Z;`);
