feat: better contextual type handling

packages/eslint-plugin/src/rules/no-unnecessary-type-assertion.ts

@@ -1,12 +1,7 @@
 import { TSESTree } from '@typescript-eslint/experimental-utils';
 import {
-  isCallExpression,
-  isJsxExpression,
-  isNewExpression,
   isObjectType,
   isObjectFlagSet,
-  isParameterDeclaration,
-  isPropertyDeclaration,
   isStrictCompilerOptionEnabled,
   isTypeFlagSet,
   isVariableDeclaration,
@@ -91,48 +86,6 @@ export default util.createRule<Options, MessageIds>({
       return true;
     }
 
-    /**
-     * Returns the contextual type of a given node.
-     * Contextual type is the type of the target the node is going into.
-     * i.e. the type of a called function's parameter, or the defined type of a variable declaration
-     */
-    function getContextualType(
-      checker: ts.TypeChecker,
-      node: ts.Expression,
-    ): ts.Type | undefined {
-      const parent = node.parent;
-      if (!parent) {
-        return;
-      }
-
-      if (isCallExpression(parent) || isNewExpression(parent)) {
-        if (node === parent.expression) {
-          // is the callee, so has no contextual type
-          return;
-        }
-      } else if (
-        isVariableDeclaration(parent) ||
-        isPropertyDeclaration(parent) ||
-        isParameterDeclaration(parent)
-      ) {
-        return parent.type
-          ? checker.getTypeFromTypeNode(parent.type)
-          : undefined;
-      } else if (isJsxExpression(parent)) {
-        return checker.getContextualType(parent);
-      } else if (
-        ![ts.SyntaxKind.TemplateSpan, ts.SyntaxKind.JsxExpression].includes(
-          parent.kind,
-        )
-      ) {
-        // parent is not something we know we can get the contextual type of
-        return;
-      }
-      // TODO - support return statement checking
-
-      return checker.getContextualType(node);
-    }
-
     /**
      * Returns true if there's a chance the variable has been used before a value has been assigned to it
      */
@@ -196,7 +149,7 @@ export default util.createRule<Options, MessageIds>({
           // we know it's a nullable type
           // so figure out if the variable is used in a place that accepts nullable types
 
-          const contextualType = getContextualType(checker, originalNode);
+          const contextualType = util.getContextualType(checker, originalNode);
           if (contextualType) {
             // in strict mode you can't assign null to undefined, so we have to make sure that
             // the two types share a nullable type

packages/eslint-plugin/src/rules/no-unsafe-return.ts

@@ -1,8 +1,9 @@
-import * as util from '../util';
 import {
   TSESTree,
   AST_NODE_TYPES,
 } from '@typescript-eslint/experimental-utils';
+import { isExpression } from 'tsutils';
+import * as util from '../util';
 
 export default util.createRule({
   name: 'no-unsafe-return',
@@ -66,17 +67,26 @@ export default util.createRule({
       }
 
       const functionNode = getParentFunctionNode(returnNode);
-      if (!functionNode?.returnType) {
+      if (!functionNode) {
         return;
       }
 
       // function has an explicit return type, so ensure it's a safe return
       const returnNodeType = checker.getTypeAtLocation(
         esTreeNodeToTSNodeMap.get(returnNode),
       );
-      const functionType = checker.getTypeAtLocation(
-        esTreeNodeToTSNodeMap.get(functionNode),
-      );
+      const functionTSNode = esTreeNodeToTSNodeMap.get(functionNode);
+
+      // function expressions will not have their return type modified based on receiver typing
+      // so we have to use the contextual typing in these cases, i.e.
+      // const foo1: () => Set<string> = () => new Set<any>();
+      // the return type of the arrow function is Set<any> even though the variable is typed as Set<string>
+      let functionType = isExpression(functionTSNode)
+        ? util.getContextualType(checker, functionTSNode)
+        : checker.getTypeAtLocation(functionTSNode);
+      if (!functionType) {
+        functionType = checker.getTypeAtLocation(functionTSNode);
+      }
 
       for (const signature of functionType.getCallSignatures()) {
         const functionReturnType = signature.getReturnType();

packages/eslint-plugin/src/util/types.ts

@@ -1,6 +1,12 @@
 import {
+  isCallExpression,
+  isJsxExpression,
+  isNewExpression,
+  isParameterDeclaration,
+  isPropertyDeclaration,
   isTypeReference,
   isUnionOrIntersectionType,
+  isVariableDeclaration,
   unionTypeParts,
 } from 'tsutils';
 import * as ts from 'typescript';
@@ -427,3 +433,43 @@ export function isUnsafeAssignment(
   }
   return false;
 }
+
+/**
+ * Returns the contextual type of a given node.
+ * Contextual type is the type of the target the node is going into.
+ * i.e. the type of a called function's parameter, or the defined type of a variable declaration
+ */
+export function getContextualType(
+  checker: ts.TypeChecker,
+  node: ts.Expression,
+): ts.Type | undefined {
+  const parent = node.parent;
+  if (!parent) {
+    return;
+  }
+
+  if (isCallExpression(parent) || isNewExpression(parent)) {
+    if (node === parent.expression) {
+      // is the callee, so has no contextual type
+      return;
+    }
+  } else if (
+    isVariableDeclaration(parent) ||
+    isPropertyDeclaration(parent) ||
+    isParameterDeclaration(parent)
+  ) {
+    return parent.type ? checker.getTypeFromTypeNode(parent.type) : undefined;
+  } else if (isJsxExpression(parent)) {
+    return checker.getContextualType(parent);
+  } else if (
+    ![ts.SyntaxKind.TemplateSpan, ts.SyntaxKind.JsxExpression].includes(
+      parent.kind,
+    )
+  ) {
+    // parent is not something we know we can get the contextual type of
+    return;
+  }
+  // TODO - support return statement checking
+
+  return checker.getContextualType(node);
+}

packages/eslint-plugin/tests/rules/no-unsafe-return.test.ts

@@ -167,5 +167,56 @@ function foo(): Set<Set<Set<string>>> { return new Set<Set<Set<any>>>(); }
         },
       ],
     }),
+    {
+      code: `
+        type Fn = () => Set<string>;
+        const foo1: Fn = () => new Set<any>();
+        const foo2: Fn = function test() { return new Set<any>() };
+      `,
+      errors: [
+        {
+          messageId: 'unsafeReturnAssignment',
+          line: 3,
+          data: {
+            sender: 'Set<any>',
+            receiver: 'Set<string>',
+          },
+        },
+        {
+          messageId: 'unsafeReturnAssignment',
+          line: 4,
+          data: {
+            sender: 'Set<any>',
+            receiver: 'Set<string>',
+          },
+        },
+      ],
+    },
+    {
+      code: `
+        type Fn = () => Set<string>;
+        function receiver(arg: Fn) {}
+        receiver(() => new Set<any>());
+        receiver(function test() { return new Set<any>() });
+      `,
+      errors: [
+        {
+          messageId: 'unsafeReturnAssignment',
+          line: 4,
+          data: {
+            sender: 'Set<any>',
+            receiver: 'Set<string>',
+          },
+        },
+        {
+          messageId: 'unsafeReturnAssignment',
+          line: 5,
+          data: {
+            sender: 'Set<any>',
+            receiver: 'Set<string>',
+          },
+        },
+      ],
+    },
   ],
 });