Fixed quick fixes for inferred type predicates (#58958)

microsoft · Jul 26, 2024 · 451388c · 451388c
1 parent 1da9630
commit 451388c
Show file tree

Hide file tree

Showing 7 changed files with 120 additions and 29 deletions.
diff --git a/src/compiler/checker.ts b/src/compiler/checker.ts
@@ -1653,6 +1653,7 @@ export function createTypeChecker(host: TypeCheckerHost): TypeChecker {
         getNonOptionalType: removeOptionalTypeMarker,
         getTypeArguments,
         typeToTypeNode: nodeBuilder.typeToTypeNode,
+        typePredicateToTypePredicateNode: nodeBuilder.typePredicateToTypePredicateNode,
         indexInfoToIndexSignatureDeclaration: nodeBuilder.indexInfoToIndexSignatureDeclaration,
         signatureToSignatureDeclaration: nodeBuilder.signatureToSignatureDeclaration,
         symbolToEntityName: nodeBuilder.symbolToEntityName,

diff --git a/src/compiler/types.ts b/src/compiler/types.ts
@@ -5074,6 +5074,7 @@ export interface TypeChecker {
     /** Note that the resulting nodes cannot be checked. */
     typeToTypeNode(type: Type, enclosingDeclaration: Node | undefined, flags: NodeBuilderFlags | undefined): TypeNode | undefined;
     /** @internal */ typeToTypeNode(type: Type, enclosingDeclaration: Node | undefined, flags: NodeBuilderFlags | undefined, tracker?: SymbolTracker): TypeNode | undefined; // eslint-disable-line @typescript-eslint/unified-signatures
+    /** @internal */ typePredicateToTypePredicateNode(typePredicate: TypePredicate, enclosingDeclaration: Node | undefined, flags: NodeBuilderFlags | undefined, tracker?: SymbolTracker): TypePredicateNode | undefined;
     /** Note that the resulting nodes cannot be checked. */
     signatureToSignatureDeclaration(signature: Signature, kind: SyntaxKind, enclosingDeclaration: Node | undefined, flags: NodeBuilderFlags | undefined): SignatureDeclaration & { typeArguments?: NodeArray<TypeNode>; } | undefined;
     /** @internal */ signatureToSignatureDeclaration(signature: Signature, kind: SyntaxKind, enclosingDeclaration: Node | undefined, flags: NodeBuilderFlags | undefined, tracker?: SymbolTracker): SignatureDeclaration & { typeArguments?: NodeArray<TypeNode>; } | undefined; // eslint-disable-line @typescript-eslint/unified-signatures

diff --git a/src/services/codefixes/fixMissingTypeAnnotationOnExports.ts b/src/services/codefixes/fixMissingTypeAnnotationOnExports.ts
@@ -91,6 +91,7 @@ import {
     TypeChecker,
     TypeFlags,
     TypeNode,
+    TypePredicate,
     UnionReduction,
     VariableDeclaration,
     VariableStatement,
@@ -103,6 +104,7 @@ import {
     createImportAdder,
     eachDiagnostic,
     registerCodeFix,
+    typePredicateToAutoImportableTypeNode,
     typeToAutoImportableTypeNode,
 } from "../_namespaces/ts.codefix.js";
 import { getIdentifierForNode } from "../refactors/helpers.js";
@@ -872,9 +874,28 @@ function withContext<T>(
             return relativeType(node);
         }
 
-        let type = isValueSignatureDeclaration(node) ?
-            tryGetReturnType(node) :
-            typeChecker.getTypeAtLocation(node);
+        let type: Type | undefined;
+
+        if (isValueSignatureDeclaration(node)) {
+            const signature = typeChecker.getSignatureFromDeclaration(node);
+            if (signature) {
+                const typePredicate = typeChecker.getTypePredicateOfSignature(signature);
+                if (typePredicate) {
+                    if (!typePredicate.type) {
+                        return emptyInferenceResult;
+                    }
+                    return {
+                        typeNode: typePredicateToTypeNode(typePredicate, findAncestor(node, isDeclaration) ?? sourceFile, getFlags(typePredicate.type)),
+                        mutatedTarget: false,
+                    };
+                }
+                type = typeChecker.getReturnTypeOfSignature(signature);
+            }
+        }
+        else {
+            type = typeChecker.getTypeAtLocation(node);
+        }
+
         if (!type) {
             return emptyInferenceResult;
         }
@@ -895,15 +916,18 @@ function withContext<T>(
         if (isParameter(node) && typeChecker.requiresAddingImplicitUndefined(node)) {
             type = typeChecker.getUnionType([typeChecker.getUndefinedType(), type], UnionReduction.None);
         }
-        const flags = (
-                isVariableDeclaration(node) ||
-                (isPropertyDeclaration(node) && hasSyntacticModifier(node, ModifierFlags.Static | ModifierFlags.Readonly))
-            ) && type.flags & TypeFlags.UniqueESSymbol ?
-            NodeBuilderFlags.AllowUniqueESSymbolType : NodeBuilderFlags.None;
         return {
-            typeNode: typeToTypeNode(type, findAncestor(node, isDeclaration) ?? sourceFile, flags),
+            typeNode: typeToTypeNode(type, findAncestor(node, isDeclaration) ?? sourceFile, getFlags(type)),
             mutatedTarget: false,
         };
+
+        function getFlags(type: Type) {
+            return (
+                    isVariableDeclaration(node) ||
+                    (isPropertyDeclaration(node) && hasSyntacticModifier(node, ModifierFlags.Static | ModifierFlags.Readonly))
+                ) && type.flags & TypeFlags.UniqueESSymbol ?
+                NodeBuilderFlags.AllowUniqueESSymbolType : NodeBuilderFlags.None;
+        }
     }
 
     function createTypeOfFromEntityNameExpression(node: EntityNameExpression) {
@@ -1084,11 +1108,18 @@ function withContext<T>(
         return isTruncated ? factory.createKeywordTypeNode(SyntaxKind.AnyKeyword) : result;
     }
 
-    function tryGetReturnType(node: SignatureDeclaration): Type | undefined {
-        const signature = typeChecker.getSignatureFromDeclaration(node);
-        if (signature) {
-            return typeChecker.getReturnTypeOfSignature(signature);
-        }
+    function typePredicateToTypeNode(typePredicate: TypePredicate, enclosingDeclaration: Node, flags = NodeBuilderFlags.None): TypeNode | undefined {
+        let isTruncated = false;
+        const result = typePredicateToAutoImportableTypeNode(typeChecker, importAdder, typePredicate, enclosingDeclaration, scriptTarget, declarationEmitNodeBuilderFlags | flags, {
+            moduleResolverHost: program,
+            trackSymbol() {
+                return true;
+            },
+            reportTruncationError() {
+                isTruncated = true;
+            },
+        });
+        return isTruncated ? factory.createKeywordTypeNode(SyntaxKind.AnyKeyword) : result;
     }
 
     function addTypeToVariableLike(decl: ParameterDeclaration | VariableDeclaration | PropertyDeclaration): DiagnosticOrDiagnosticAndArguments | undefined {

diff --git a/src/services/codefixes/helpers.ts b/src/services/codefixes/helpers.ts
@@ -105,6 +105,7 @@ import {
     TypeFlags,
     TypeNode,
     TypeParameterDeclaration,
+    TypePredicate,
     unescapeLeadingUnderscores,
     UnionType,
     UserPreferences,
@@ -602,6 +603,20 @@ export function typeToAutoImportableTypeNode(checker: TypeChecker, importAdder:
     return getSynthesizedDeepClone(typeNode);
 }
 
+/** @internal */
+export function typePredicateToAutoImportableTypeNode(checker: TypeChecker, importAdder: ImportAdder, typePredicate: TypePredicate, contextNode: Node | undefined, scriptTarget: ScriptTarget, flags?: NodeBuilderFlags, tracker?: SymbolTracker): TypeNode | undefined {
+    let typePredicateNode = checker.typePredicateToTypePredicateNode(typePredicate, contextNode, flags, tracker);
+    if (typePredicateNode?.type && isImportTypeNode(typePredicateNode.type)) {
+        const importableReference = tryGetAutoImportableReferenceFromTypeNode(typePredicateNode.type, scriptTarget);
+        if (importableReference) {
+            importSymbols(importAdder, importableReference.symbols);
+            typePredicateNode = factory.updateTypePredicateNode(typePredicateNode, typePredicateNode.assertsModifier, typePredicateNode.parameterName, importableReference.typeNode);
+        }
+    }
+    // Ensure nodes are fresh so they can have different positions when going through formatting.
+    return getSynthesizedDeepClone(typePredicateNode);
+}
+
 function typeContainsTypeParameter(type: Type) {
     if (type.isUnionOrIntersection()) {
         return type.types.some(typeContainsTypeParameter);

diff --git a/src/services/refactors/inferFunctionReturnType.ts b/src/services/refactors/inferFunctionReturnType.ts
@@ -24,7 +24,6 @@ import {
     SyntaxKind,
     textChanges,
     Type,
-    TypeChecker,
     TypeNode,
 } from "../_namespaces/ts.js";
 import {
@@ -113,7 +112,31 @@ function getInfo(context: RefactorContext): FunctionInfo | RefactorErrorInfo | u
     }
 
     const typeChecker = context.program.getTypeChecker();
-    const returnType = tryGetReturnType(typeChecker, declaration);
+
+    let returnType: Type | undefined;
+
+    if (typeChecker.isImplementationOfOverload(declaration)) {
+        const signatures = typeChecker.getTypeAtLocation(declaration).getCallSignatures();
+        if (signatures.length > 1) {
+            returnType = typeChecker.getUnionType(mapDefined(signatures, s => s.getReturnType()));
+        }
+    }
+    if (!returnType) {
+        const signature = typeChecker.getSignatureFromDeclaration(declaration);
+        if (signature) {
+            const typePredicate = typeChecker.getTypePredicateOfSignature(signature);
+            if (typePredicate && typePredicate.type) {
+                const typePredicateTypeNode = typeChecker.typePredicateToTypePredicateNode(typePredicate, declaration, NodeBuilderFlags.NoTruncation);
+                if (typePredicateTypeNode) {
+                    return { declaration, returnTypeNode: typePredicateTypeNode };
+                }
+            }
+            else {
+                returnType = typeChecker.getReturnTypeOfSignature(signature);
+            }
+        }
+    }
+
     if (!returnType) {
         return { error: getLocaleSpecificMessage(Diagnostics.Could_not_determine_function_return_type) };
     }
@@ -135,16 +158,3 @@ function isConvertibleDeclaration(node: Node): node is ConvertibleDeclaration {
             return false;
     }
 }
-
-function tryGetReturnType(typeChecker: TypeChecker, node: ConvertibleDeclaration): Type | undefined {
-    if (typeChecker.isImplementationOfOverload(node)) {
-        const signatures = typeChecker.getTypeAtLocation(node).getCallSignatures();
-        if (signatures.length > 1) {
-            return typeChecker.getUnionType(mapDefined(signatures, s => s.getReturnType()));
-        }
-    }
-    const signature = typeChecker.getSignatureFromDeclaration(node);
-    if (signature) {
-        return typeChecker.getReturnTypeOfSignature(signature);
-    }
-}
diff --git a/tests/cases/fourslash/codeFixMissingTypeAnnotationOnExportsTypePredicate1.ts b/tests/cases/fourslash/codeFixMissingTypeAnnotationOnExportsTypePredicate1.ts
@@ -0,0 +1,17 @@
+/// <reference path='fourslash.ts'/>
+
+// @isolatedDeclarations: true
+// @declaration: true
+
+// @filename: index.ts
+//// export function isString(value: unknown) {
+////   return typeof value === "string";
+//// }
+
+verify.codeFix({
+    description: `Add return type 'value is string'`,
+    index: 0,
+    newFileContent: `export function isString(value: unknown): value is string {
+  return typeof value === "string";
+}`,
+});
diff --git a/tests/cases/fourslash/refactorInferFunctionReturnTypePredicate1.ts b/tests/cases/fourslash/refactorInferFunctionReturnTypePredicate1.ts
@@ -0,0 +1,16 @@
+/// <reference path='fourslash.ts' />
+
+//// function /*a*//*b*/isString(value: unknown) {
+////   return typeof value === "string";
+//// }
+
+goTo.select("a", "b");
+edit.applyRefactor({
+    refactorName: "Infer function return type",
+    actionName: "Infer function return type",
+    actionDescription: "Infer function return type",
+    newContent:
+`function isString(value: unknown): value is string {
+  return typeof value === "string";
+}`
+});