From 8ba501926a48c244b13ef3c96c2bf0942628e003 Mon Sep 17 00:00:00 2001
From: Wenlu Wang <kingwenlu@gmail.com>
Date: Fri, 24 Aug 2018 12:31:53 +0800
Subject: [PATCH] try get add missing member return type from context (#26250)

* try get add missing member return type from context

* support contextual type
---
 src/services/codefixes/helpers.ts             | 22 ++++++++-----------
 .../fourslash/codeFixAddMissingMember10.ts    | 16 ++++++++++++++
 .../fourslash/codeFixAddMissingMember11.ts    | 18 +++++++++++++++
 3 files changed, 43 insertions(+), 13 deletions(-)
 create mode 100644 tests/cases/fourslash/codeFixAddMissingMember10.ts
 create mode 100644 tests/cases/fourslash/codeFixAddMissingMember11.ts

diff --git a/src/services/codefixes/helpers.ts b/src/services/codefixes/helpers.ts
index 5b955d2492..eec693bbb6 100644
--- a/src/services/codefixes/helpers.ts
+++ b/src/services/codefixes/helpers.ts
@@ -112,27 +112,23 @@ namespace ts.codefix {
 
     export function createMethodFromCallExpression(
         context: CodeFixContextBase,
-        { typeArguments, arguments: args, parent: parent }: CallExpression,
+        call: CallExpression,
         methodName: string,
         inJs: boolean,
         makeStatic: boolean,
         preferences: UserPreferences,
         body: boolean,
     ): MethodDeclaration {
+        const { typeArguments, arguments: args, parent } = call;
         const checker = context.program.getTypeChecker();
-        const types = map(args,
-            arg => {
-                let type = checker.getTypeAtLocation(arg);
-                if (type === undefined) {
-                    return undefined;
-                }
-                // Widen the type so we don't emit nonsense annotations like "function fn(x: 3) {"
-                type = checker.getBaseTypeOfLiteralType(type);
-                return checker.typeToTypeNode(type);
-            });
+        const types = map(args, arg =>
+            // Widen the type so we don't emit nonsense annotations like "function fn(x: 3) {"
+            checker.typeToTypeNode(checker.getBaseTypeOfLiteralType(checker.getTypeAtLocation(arg))));
         const names = map(args, arg =>
             isIdentifier(arg) ? arg.text :
-            isPropertyAccessExpression(arg) ? arg.name.text : undefined);
+                isPropertyAccessExpression(arg) ? arg.name.text : undefined);
+        const contextualType = checker.getContextualType(call);
+        const returnType = inJs ? undefined : contextualType && checker.typeToTypeNode(contextualType, call) || createKeywordTypeNode(SyntaxKind.AnyKeyword);
         return createMethod(
             /*decorators*/ undefined,
             /*modifiers*/ makeStatic ? [createToken(SyntaxKind.StaticKeyword)] : undefined,
@@ -142,7 +138,7 @@ namespace ts.codefix {
             /*typeParameters*/ inJs ? undefined : map(typeArguments, (_, i) =>
                 createTypeParameterDeclaration(CharacterCodes.T + typeArguments!.length - 1 <= CharacterCodes.Z ? String.fromCharCode(CharacterCodes.T + i) : `T${i}`)),
             /*parameters*/ createDummyParameters(args.length, names, types, /*minArgumentCount*/ undefined, inJs),
-            /*type*/ inJs ? undefined : createKeywordTypeNode(SyntaxKind.AnyKeyword),
+            /*type*/ returnType,
             body ? createStubbedMethodBody(preferences) : undefined);
     }
 
diff --git a/tests/cases/fourslash/codeFixAddMissingMember10.ts b/tests/cases/fourslash/codeFixAddMissingMember10.ts
new file mode 100644
index 0000000000..990d86ec49
--- /dev/null
+++ b/tests/cases/fourslash/codeFixAddMissingMember10.ts
@@ -0,0 +1,16 @@
+/// <reference path='fourslash.ts' />
+
+//// class C {}
+//// const n: number = new C().add(1, 2);
+
+verify.codeFixAll({
+    fixId: "addMissingMember",
+    fixAllDescription: "Add all missing members",
+    newFileContent:
+`class C {
+    add(arg0: number, arg1: number): number {
+        throw new Error("Method not implemented.");
+    }
+}
+const n: number = new C().add(1, 2);`,
+});
diff --git a/tests/cases/fourslash/codeFixAddMissingMember11.ts b/tests/cases/fourslash/codeFixAddMissingMember11.ts
new file mode 100644
index 0000000000..6f68502d4c
--- /dev/null
+++ b/tests/cases/fourslash/codeFixAddMissingMember11.ts
@@ -0,0 +1,18 @@
+/// <reference path='fourslash.ts' />
+
+//// class C {}
+//// function f(v: number): void { }
+//// f(new C().add(1, 2))
+
+verify.codeFixAll({
+    fixId: "addMissingMember",
+    fixAllDescription: "Add all missing members",
+    newFileContent:
+`class C {
+    add(arg0: number, arg1: number): number {
+        throw new Error("Method not implemented.");
+    }
+}
+function f(v: number): void { }
+f(new C().add(1, 2))`,
+});