Fix codefix-triggered auto import of aliased exports

src/services/codefixes/fixAddMissingParam.ts

@@ -350,7 +350,7 @@ function isOptionalPos(declarations: ConvertibleSignatureDeclaration[], pos: num
 function getParameterType(importAdder: ImportAdder, typeNode: TypeNode | undefined, scriptTarget: ScriptTarget) {
     const importableReference = tryGetAutoImportableReferenceFromTypeNode(typeNode, scriptTarget);
     if (importableReference) {
-        importSymbols(importAdder, importableReference.symbols);
+        importSymbols(importAdder, importableReference.replacements);
         return importableReference.typeNode;
     }
     return typeNode;

src/services/codefixes/helpers.ts

@@ -230,7 +230,7 @@ export function addNewNodeForMemberSymbol(
                 const importableReference = tryGetAutoImportableReferenceFromTypeNode(typeNode, scriptTarget);
                 if (importableReference) {
                     typeNode = importableReference.typeNode;
-                    importSymbols(importAdder, importableReference.symbols);
+                    importSymbols(importAdder, importableReference.replacements);
                 }
             }
             addClassElement(factory.createPropertyDeclaration(
@@ -253,7 +253,7 @@ export function addNewNodeForMemberSymbol(
                 const importableReference = tryGetAutoImportableReferenceFromTypeNode(typeNode, scriptTarget);
                 if (importableReference) {
                     typeNode = importableReference.typeNode;
-                    importSymbols(importAdder, importableReference.symbols);
+                    importSymbols(importAdder, importableReference.replacements);
                 }
             }
             for (const accessor of orderedAccessors) {
@@ -415,14 +415,14 @@ export function createSignatureDeclarationFromSignature(
                     const importableReference = tryGetAutoImportableReferenceFromTypeNode(constraint, scriptTarget);
                     if (importableReference) {
                         constraint = importableReference.typeNode;
-                        importSymbols(importAdder, importableReference.symbols);
+                        importSymbols(importAdder, importableReference.replacements);
                     }
                 }
                 if (defaultType) {
                     const importableReference = tryGetAutoImportableReferenceFromTypeNode(defaultType, scriptTarget);
                     if (importableReference) {
                         defaultType = importableReference.typeNode;
-                        importSymbols(importAdder, importableReference.symbols);
+                        importSymbols(importAdder, importableReference.replacements);
                     }
                 }
                 return factory.updateTypeParameterDeclaration(
@@ -443,7 +443,7 @@ export function createSignatureDeclarationFromSignature(
                 const importableReference = tryGetAutoImportableReferenceFromTypeNode(type, scriptTarget);
                 if (importableReference) {
                     type = importableReference.typeNode;
-                    importSymbols(importAdder, importableReference.symbols);
+                    importSymbols(importAdder, importableReference.replacements);
                 }
             }
             return factory.updateParameterDeclaration(
@@ -463,7 +463,7 @@ export function createSignatureDeclarationFromSignature(
             const importableReference = tryGetAutoImportableReferenceFromTypeNode(type, scriptTarget);
             if (importableReference) {
                 type = importableReference.typeNode;
-                importSymbols(importAdder, importableReference.symbols);
+                importSymbols(importAdder, importableReference.replacements);
             }
         }
     }
@@ -609,7 +609,7 @@ export function typeNodeToAutoImportableTypeNode(typeNode: TypeNode, importAdder
     if (typeNode && isImportTypeNode(typeNode)) {
         const importableReference = tryGetAutoImportableReferenceFromTypeNode(typeNode, scriptTarget);
         if (importableReference) {
-            importSymbols(importAdder, importableReference.symbols);
+            importSymbols(importAdder, importableReference.replacements);
             typeNode = importableReference.typeNode;
         }
     }
@@ -658,7 +658,7 @@ export function typePredicateToAutoImportableTypeNode(checker: TypeChecker, impo
     if (typePredicateNode?.type && isImportTypeNode(typePredicateNode.type)) {
         const importableReference = tryGetAutoImportableReferenceFromTypeNode(typePredicateNode.type, scriptTarget);
         if (importableReference) {
-            importSymbols(importAdder, importableReference.symbols);
+            importSymbols(importAdder, importableReference.replacements);
             typePredicateNode = factory.updateTypePredicateNode(typePredicateNode, typePredicateNode.assertsModifier, typePredicateNode.parameterName, importableReference.typeNode);
         }
     }
@@ -945,6 +945,19 @@ function findJsonProperty(obj: ObjectLiteralExpression, name: string): PropertyA
     return find(obj.properties, (p): p is PropertyAssignment => isPropertyAssignment(p) && !!p.name && isStringLiteral(p.name) && p.name.text === name);
 }
 
+/**
+ * @internal
+ * Information about an `import("...").SomeType` replaced by an auto-importable type reference.
+ */
+export interface ImportTypeReplacement {
+    symbol: Symbol;
+    /**
+     * The text of the new identifier generated to reference the symbol. May differ from
+     * `symbol.name` for default exports and aliased exports like `export { Foo as Bar }`.
+     */
+    referenceName: string;
+}
+
 /**
  * Given a type node containing 'import("./a").SomeType<import("./b").OtherType<...>>',
  * returns an equivalent type reference node with any nested ImportTypeNodes also replaced
@@ -954,12 +967,12 @@ function findJsonProperty(obj: ObjectLiteralExpression, name: string): PropertyA
  */
 export function tryGetAutoImportableReferenceFromTypeNode(importTypeNode: TypeNode | undefined, scriptTarget: ScriptTarget): {
     typeNode: TypeNode;
-    symbols: Symbol[];
+    replacements: ImportTypeReplacement[];
 } | undefined {
-    let symbols: Symbol[] | undefined;
+    let replacements: ImportTypeReplacement[] | undefined;
     const typeNode = visitNode(importTypeNode, visit, isTypeNode);
-    if (symbols && typeNode) {
-        return { typeNode, symbols };
+    if (replacements && typeNode) {
+        return { typeNode, replacements };
     }
 
     function visit(node: Node): Node {
@@ -972,12 +985,14 @@ export function tryGetAutoImportableReferenceFromTypeNode(importTypeNode: TypeNo
                 // it can't refer to reserved internal symbol names and such
                 return visitEachChild(node, visit, /*context*/ undefined);
             }
-            const name = getNameForExportedSymbol(firstIdentifier.symbol, scriptTarget);
-            const qualifier = name !== firstIdentifier.text
-                ? replaceFirstIdentifierOfEntityName(node.qualifier, factory.createIdentifier(name))
+            const referenceName = firstIdentifier.text === "default"
+                ? getNameForExportedSymbol(firstIdentifier.symbol, scriptTarget)
+                : firstIdentifier.text;
+            const qualifier = referenceName !== firstIdentifier.text
+                ? replaceFirstIdentifierOfEntityName(node.qualifier, factory.createIdentifier(referenceName))
                 : node.qualifier;
 
-            symbols = append(symbols, firstIdentifier.symbol);
+            replacements = append(replacements, { symbol: firstIdentifier.symbol, referenceName });
             const typeArguments = visitNodes(node.typeArguments, visit, isTypeNode);
             return factory.createTypeReferenceNode(qualifier, typeArguments);
         }
@@ -993,8 +1008,8 @@ function replaceFirstIdentifierOfEntityName(name: EntityName, newIdentifier: Ide
 }
 
 /** @internal */
-export function importSymbols(importAdder: ImportAdder, symbols: readonly Symbol[]): void {
-    symbols.forEach(s => importAdder.addImportFromExportedSymbol(s, /*isValidTypeOnlyUseSite*/ true));
+export function importSymbols(importAdder: ImportAdder, symbols: readonly ImportTypeReplacement[]): void {
+    symbols.forEach(importAdder.addImportFromImportTypeReplacement);
 }
 
 /** @internal */

src/services/codefixes/importFixes.ts

@@ -2,6 +2,7 @@ import {
     createCodeFixAction,
     createCombinedCodeActions,
     eachDiagnostic,
+    ImportTypeReplacement,
     registerCodeFix,
 } from "../_namespaces/ts.codefix.js";
 import {
@@ -227,6 +228,7 @@ export interface ImportAdder {
     hasFixes(): boolean;
     addImportFromDiagnostic: (diagnostic: DiagnosticWithLocation, context: CodeFixContextBase) => void;
     addImportFromExportedSymbol: (exportedSymbol: Symbol, isValidTypeOnlyUseSite?: boolean, referenceImport?: ImportOrRequireAliasDeclaration) => void;
+    addImportFromImportTypeReplacement: (replacement: ImportTypeReplacement) => void;
     addImportForNonExistentExport: (exportName: string, exportingFileName: string, exportKind: ExportKind, exportedMeanings: SymbolFlags, isImportUsageValidAsTypeOnly: boolean) => void;
     addImportForUnresolvedIdentifier: (context: CodeFixContextBase, symbolToken: Identifier, useAutoImportProvider: boolean) => void;
     addVerbatimImport: (declaration: AnyImportOrRequireStatement | ImportOrRequireAliasDeclaration) => void;
@@ -257,7 +259,7 @@ function createImportAdderWorker(sourceFile: SourceFile | FutureSourceFile, prog
     type NewImportsKey = `${0 | 1}|${string}`;
     /** Use `getNewImportEntry` for access */
     const newImports = new Map<NewImportsKey, Mutable<ImportsCollection & { useRequire: boolean; }>>();
-    return { addImportFromDiagnostic, addImportFromExportedSymbol, writeFixes, hasFixes, addImportForUnresolvedIdentifier, addImportForNonExistentExport, removeExistingImport, addVerbatimImport };
+    return { addImportFromDiagnostic, addImportFromExportedSymbol, addImportFromImportTypeReplacement, writeFixes, hasFixes, addImportForUnresolvedIdentifier, addImportForNonExistentExport, removeExistingImport, addVerbatimImport };
 
     function addVerbatimImport(declaration: AnyImportOrRequireStatement | ImportOrRequireAliasDeclaration) {
         verbatimImports.add(declaration);
@@ -317,6 +319,23 @@ function createImportAdderWorker(sourceFile: SourceFile | FutureSourceFile, prog
         }
     }
 
+    function addImportFromImportTypeReplacement(replacement: ImportTypeReplacement) {
+        const checker = program.getTypeChecker();
+        const symbol = checker.getMergedSymbol(skipAlias(replacement.symbol, checker));
+        const moduleSymbol = Debug.checkDefined(replacement.symbol.parent);
+        const exportInfo = getAllExportInfoForSymbol(sourceFile, symbol, replacement.referenceName, moduleSymbol, /*preferCapitalized*/ false, program, host, preferences, cancellationToken);
+        if (!exportInfo) {
+            // If no exportInfo is found, this means export could not be resolved when we have filtered for autoImportFileExcludePatterns,
+            //     so we should not generate an import.
+            Debug.assert(preferences.autoImportFileExcludePatterns?.length);
+            return;
+        }
+        const fix = getImportFixForSymbol(sourceFile, exportInfo, program, /*position*/ undefined, /*isValidTypeOnlyUseSite*/ true, /*useRequire*/ false, host, preferences);
+        if (fix) {
+            addImport({ fix, symbolName: replacement.referenceName, errorIdentifierText: undefined });
+        }
+    }
+
     function addImportForNonExistentExport(exportName: string, exportingFileName: string, exportKind: ExportKind, exportedMeanings: SymbolFlags, isImportUsageValidAsTypeOnly: boolean) {
         const exportingSourceFile = program.getSourceFile(exportingFileName);
         const useRequire = shouldUseRequire(sourceFile, program);

src/services/codefixes/inferFromUsage.ts

@@ -439,7 +439,7 @@ function tryReplaceImportTypeNodeWithAutoImport(
 ): boolean {
     const importableReference = tryGetAutoImportableReferenceFromTypeNode(typeNode, scriptTarget);
     if (importableReference && changes.tryInsertTypeAnnotation(sourceFile, declaration, importableReference.typeNode)) {
-        forEach(importableReference.symbols, s => importAdder.addImportFromExportedSymbol(s, /*isValidTypeOnlyUseSite*/ true));
+        forEach(importableReference.replacements, importAdder.addImportFromImportTypeReplacement);
         return true;
     }
     return false;

tests/cases/fourslash/autoImportAliasedModuleAugmentation.ts

@@ -0,0 +1,58 @@
+/// <reference path="fourslash.ts" />
+
+// @module: nodenext
+
+// @Filename: /node_modules/@sapphire/pieces/index.d.ts
+//// interface Container {
+////   stores: unknown;
+//// }
+////
+//// declare class Piece {
+////   container: Container;
+//// }
+////
+//// export { Piece, type Container as Alias };
+
+// @FileName: /augmentation.ts
+//// declare module "@sapphire/pieces" {
+////   interface Alias {
+////     client: unknown;
+////   }
+//// }
+
+// @Filename: /index.ts
+//// import { Piece } from "@sapphire/pieces";
+//// class FullPiece extends Piece {
+////   /*1*/
+//// }
+
+const preferences = {
+  includeCompletionsWithClassMemberSnippets: true,
+  includeCompletionsWithInsertText: true,
+};
+
+verify.completions({
+  marker: "1",
+  includes: [
+    {
+      name: "container",
+      insertText: "container: Alias;",
+      filterText: "container",
+      hasAction: true,
+      source: "ClassMemberSnippet/",
+    },
+  ],
+  preferences,
+  isNewIdentifierLocation: true,
+});
+
+verify.applyCodeActionFromCompletion("1", {
+  name: "container",
+  source: "ClassMemberSnippet/",
+  description: `Includes imports of types referenced by 'container'`,
+  newFileContent: `import { Alias, Piece } from "@sapphire/pieces";
+class FullPiece extends Piece {
+
+}`,
+  preferences,
+});