feat(eslint-plugin): additional annotation spacing rules for va… (typescript-eslint#1496)

chkt · web-flow · commit b097245df351 · 2020-03-01T21:12:41.000-08:00
diff --git a/packages/eslint-plugin/docs/rules/type-annotation-spacing.md b/packages/eslint-plugin/docs/rules/type-annotation-spacing.md
@@ -41,7 +41,7 @@ This rule has an object option:
 - `"before": true`, (default for arrow) requires a space before the colon/arrow.
 - `"after": true`, (default) requires a space after the colon/arrow.
 - `"after": false`, disallows spaces after the colon/arrow.
-- `"overrides"`, overrides the default options for type annotations with `colon` (e.g. `const foo: string`) and function types with `arrow` (e.g. `type Foo = () => {}`).
+- `"overrides"`, overrides the default options for type annotations with `colon` (e.g. `const foo: string`) and function types with `arrow` (e.g. `type Foo = () => {}`). Additionally allows granular overrides for `variable` (`const foo: string`),`parameter` (`function foo(bar: string) {...}`),`property` (`interface Foo { bar: string }`) and `returnType` (`function foo(): string {...}`) annotations.
 
 ### defaults
 
diff --git a/packages/eslint-plugin/src/rules/type-annotation-spacing.ts b/packages/eslint-plugin/src/rules/type-annotation-spacing.ts
@@ -1,22 +1,35 @@
 import { TSESTree } from '@typescript-eslint/experimental-utils';
 import * as util from '../util';
+import {
+  isClassOrTypeElement,
+  isFunction,
+  isFunctionOrFunctionType,
+  isIdentifier,
+  isTSFunctionType,
+  isVariableDeclarator,
+} from '../util';
 
-type Options = [
-  {
-    before?: boolean;
-    after?: boolean;
-    overrides?: {
-      colon?: {
-        before?: boolean;
-        after?: boolean;
-      };
-      arrow?: {
-        before?: boolean;
-        after?: boolean;
-      };
-    };
-  }?,
-];
+interface WhitespaceRule {
+  readonly before?: boolean;
+  readonly after?: boolean;
+}
+
+interface WhitespaceOverride {
+  readonly colon?: WhitespaceRule;
+  readonly arrow?: WhitespaceRule;
+  readonly variable?: WhitespaceRule;
+  readonly property?: WhitespaceRule;
+  readonly parameter?: WhitespaceRule;
+  readonly returnType?: WhitespaceRule;
+}
+
+interface Config extends WhitespaceRule {
+  readonly overrides?: WhitespaceOverride;
+}
+
+type WhitespaceRules = Required<WhitespaceOverride>;
+
+type Options = [Config?];
 type MessageIds =
   | 'expectedSpaceAfter'
   | 'expectedSpaceBefore'
@@ -32,6 +45,67 @@ const definition = {
   additionalProperties: false,
 };
 
+function createRules(options?: Config): WhitespaceRules {
+  const globals = {
+    ...(options?.before !== undefined ? { before: options.before } : {}),
+    ...(options?.after !== undefined ? { after: options.after } : {}),
+  };
+  const override = options?.overrides ?? {};
+  const colon = {
+    ...{ before: false, after: true },
+    ...globals,
+    ...override?.colon,
+  };
+  const arrow = {
+    ...{ before: true, after: true },
+    ...globals,
+    ...override?.arrow,
+  };
+
+  return {
+    colon: colon,
+    arrow: arrow,
+    variable: { ...colon, ...override?.variable },
+    property: { ...colon, ...override?.property },
+    parameter: { ...colon, ...override?.parameter },
+    returnType: { ...colon, ...override?.returnType },
+  };
+}
+
+function getIdentifierRules(
+  rules: WhitespaceRules,
+  node: TSESTree.Node | undefined,
+): WhitespaceRule {
+  const scope = node?.parent;
+
+  if (isVariableDeclarator(scope)) {
+    return rules.variable;
+  } else if (isFunctionOrFunctionType(scope)) {
+    return rules.parameter;
+  } else {
+    return rules.colon;
+  }
+}
+
+function getRules(
+  rules: WhitespaceRules,
+  node: TSESTree.TypeNode,
+): WhitespaceRule {
+  const scope = node?.parent?.parent;
+
+  if (isTSFunctionType(scope)) {
+    return rules.arrow;
+  } else if (isIdentifier(scope)) {
+    return getIdentifierRules(rules, scope);
+  } else if (isClassOrTypeElement(scope)) {
+    return rules.property;
+  } else if (isFunction(scope)) {
+    return rules.returnType;
+  } else {
+    return rules.colon;
+  }
+}
+
 export default util.createRule<Options, MessageIds>({
   name: 'type-annotation-spacing',
   meta: {
@@ -59,6 +133,10 @@ export default util.createRule<Options, MessageIds>({
             properties: {
               colon: definition,
               arrow: definition,
+              variable: definition,
+              parameter: definition,
+              property: definition,
+              returnType: definition,
             },
             additionalProperties: false,
           },
@@ -76,20 +154,7 @@ export default util.createRule<Options, MessageIds>({
     const punctuators = [':', '=>'];
     const sourceCode = context.getSourceCode();
 
-    const overrides = options?.overrides ?? { colon: {}, arrow: {} };
-
-    const colonOptions = Object.assign(
-      {},
-      { before: false, after: true },
-      options,
-      overrides.colon,
-    );
-    const arrowOptions = Object.assign(
-      {},
-      { before: true, after: true },
-      options,
-      overrides.arrow,
-    );
+    const ruleSet = createRules(options);
 
     /**
      * Checks if there's proper spacing around type annotations (no space
@@ -108,8 +173,7 @@ export default util.createRule<Options, MessageIds>({
         return;
       }
 
-      const before = type === ':' ? colonOptions.before : arrowOptions.before;
-      const after = type === ':' ? colonOptions.after : arrowOptions.after;
+      const { before, after } = getRules(ruleSet, typeAnnotation);
 
       if (type === ':' && previousToken.value === '?') {
         // shift the start to the ?
diff --git a/packages/eslint-plugin/src/util/astUtils.ts b/packages/eslint-plugin/src/util/astUtils.ts
@@ -82,6 +82,95 @@ function isTypeAssertion(
   );
 }
 
+function isVariableDeclarator(
+  node: TSESTree.Node | undefined,
+): node is TSESTree.VariableDeclarator {
+  return node?.type === AST_NODE_TYPES.VariableDeclarator;
+}
+
+function isFunction(
+  node: TSESTree.Node | undefined,
+): node is
+  | TSESTree.ArrowFunctionExpression
+  | TSESTree.FunctionDeclaration
+  | TSESTree.FunctionExpression {
+  if (!node) {
+    return false;
+  }
+
+  return [
+    AST_NODE_TYPES.ArrowFunctionExpression,
+    AST_NODE_TYPES.FunctionDeclaration,
+    AST_NODE_TYPES.FunctionExpression,
+  ].includes(node.type);
+}
+
+function isFunctionType(
+  node: TSESTree.Node | undefined,
+): node is
+  | TSESTree.TSCallSignatureDeclaration
+  | TSESTree.TSConstructSignatureDeclaration
+  | TSESTree.TSEmptyBodyFunctionExpression
+  | TSESTree.TSFunctionType
+  | TSESTree.TSMethodSignature {
+  if (!node) {
+    return false;
+  }
+
+  return [
+    AST_NODE_TYPES.TSCallSignatureDeclaration,
+    AST_NODE_TYPES.TSConstructSignatureDeclaration,
+    AST_NODE_TYPES.TSEmptyBodyFunctionExpression,
+    AST_NODE_TYPES.TSFunctionType,
+    AST_NODE_TYPES.TSMethodSignature,
+  ].includes(node.type);
+}
+
+function isFunctionOrFunctionType(
+  node: TSESTree.Node | undefined,
+): node is
+  | TSESTree.ArrowFunctionExpression
+  | TSESTree.FunctionDeclaration
+  | TSESTree.FunctionExpression
+  | TSESTree.TSCallSignatureDeclaration
+  | TSESTree.TSConstructSignatureDeclaration
+  | TSESTree.TSEmptyBodyFunctionExpression
+  | TSESTree.TSFunctionType
+  | TSESTree.TSMethodSignature {
+  return isFunction(node) || isFunctionType(node);
+}
+
+function isTSFunctionType(
+  node: TSESTree.Node | undefined,
+): node is TSESTree.TSFunctionType {
+  return node?.type === AST_NODE_TYPES.TSFunctionType;
+}
+
+function isClassOrTypeElement(
+  node: TSESTree.Node | undefined,
+): node is TSESTree.ClassElement | TSESTree.TypeElement {
+  if (!node) {
+    return false;
+  }
+
+  return [
+    // ClassElement
+    AST_NODE_TYPES.ClassProperty,
+    AST_NODE_TYPES.FunctionExpression,
+    AST_NODE_TYPES.MethodDefinition,
+    AST_NODE_TYPES.TSAbstractClassProperty,
+    AST_NODE_TYPES.TSAbstractMethodDefinition,
+    AST_NODE_TYPES.TSEmptyBodyFunctionExpression,
+    AST_NODE_TYPES.TSIndexSignature,
+    // TypeElement
+    AST_NODE_TYPES.TSCallSignatureDeclaration,
+    AST_NODE_TYPES.TSConstructSignatureDeclaration,
+    // AST_NODE_TYPES.TSIndexSignature,
+    AST_NODE_TYPES.TSMethodSignature,
+    AST_NODE_TYPES.TSPropertySignature,
+  ].includes(node.type);
+}
+
 /**
  * Checks if a node is a constructor method.
  */
@@ -136,6 +225,10 @@ export {
   isAwaitExpression,
   isAwaitKeyword,
   isConstructor,
+  isClassOrTypeElement,
+  isFunction,
+  isFunctionOrFunctionType,
+  isFunctionType,
   isIdentifier,
   isLogicalOrOperator,
   isNonNullAssertionPunctuator,
@@ -145,6 +238,8 @@ export {
   isOptionalOptionalChain,
   isSetter,
   isTokenOnSameLine,
+  isTSFunctionType,
   isTypeAssertion,
+  isVariableDeclarator,
   LINEBREAK_MATCHER,
 };
diff --git a/packages/eslint-plugin/tests/rules/type-annotation-spacing.test.ts b/packages/eslint-plugin/tests/rules/type-annotation-spacing.test.ts
