Refactor to use assertion handlers for easier extension

source/index.ts

@@ -1,4 +1,4 @@
 import tsd from './lib';
 
-export * from './lib/assert';
+export * from './lib/assertions/assert';
 export default tsd;

source/lib/assertions/handlers/handler.ts

@@ -0,0 +1,12 @@
+import {TypeChecker, CallExpression} from '../../../../libraries/typescript/lib/typescript';
+import {Diagnostic} from '../../interfaces';
+
+/**
+ * A handler is a method which accepts the TypeScript type checker together with a set of assertion nodes. The type checker
+ * can be used to retrieve extra type information from these nodes in order to determine a list of diagnostics.
+ *
+ * @param typeChecker - The TypeScript type checker.
+ * @param nodes - List of nodes.
+ * @returns List of diagnostics.
+ */
+export type Handler = (typeChecker: TypeChecker, nodes: Set<CallExpression>) => Diagnostic[];

source/lib/assertions/handlers/index.ts

@@ -0,0 +1,4 @@
+export {Handler} from './handler';
+
+// Handlers
+export {strictAssertion} from './strict-assertion';

source/lib/assertions/handlers/strict-assertion.ts

@@ -0,0 +1,51 @@
+import {TypeChecker, CallExpression} from '../../../../libraries/typescript/lib/typescript';
+import {Diagnostic} from '../../interfaces';
+
+/**
+ * Performs strict type assertion between the argument if the assertion, and the generic type of the assertion.
+ *
+ * @param checker - The TypeScript type checker.
+ * @param nodes - The `expectType` AST nodes.
+ * @return List of custom diagnostics.
+ */
+export const strictAssertion = (checker: TypeChecker, nodes: Set<CallExpression>): Diagnostic[] => {
+	const diagnostics: Diagnostic[] = [];
+
+	if (!nodes) {
+		return diagnostics;
+	}
+
+	for (const node of nodes) {
+		if (!node.typeArguments) {
+			// Skip if the node does not have generics
+			continue;
+		}
+
+		// Retrieve the type to be expected. This is the type inside the generic.
+		const expectedType = checker.getTypeFromTypeNode(node.typeArguments[0]);
+		const argumentType = checker.getTypeAtLocation(node.arguments[0]);
+
+		if (!checker.isAssignableTo(argumentType, expectedType)) {
+			// The argument type is not assignable to the expected type. TypeScript will catch this for us.
+			continue;
+		}
+
+		if (!checker.isAssignableTo(expectedType, argumentType)) { // tslint:disable-line:early-exit
+			/**
+			 * At this point, the expected type is not assignable to the argument type, but the argument type is
+			 * assignable to the expected type. This means our type is too wide.
+			 */
+			const position = node.getSourceFile().getLineAndCharacterOfPosition(node.getStart());
+
+			diagnostics.push({
+				fileName: node.getSourceFile().fileName,
+				message: `Parameter type \`${checker.typeToString(expectedType)}\` is declared too wide for argument type \`${checker.typeToString(argumentType)}\`.`,
+				severity: 'error',
+				line: position.line + 1,
+				column: position.character,
+			});
+		}
+	}
+
+	return diagnostics;
+};

source/lib/assertions/index.ts

@@ -0,0 +1,42 @@
+import {TypeChecker, CallExpression} from '../../../libraries/typescript/lib/typescript';
+import {Diagnostic} from '../interfaces';
+import {Handler, strictAssertion} from './handlers';
+
+export enum Assertion {
+	EXPECT_TYPE = 'expectType',
+	EXPECT_ERROR = 'expectError'
+}
+
+// List of diagnostic handlers attached to the assertion
+const assertionHandlers = new Map<string, Handler | Handler[]>([
+	[Assertion.EXPECT_TYPE, strictAssertion]
+]);
+
+/**
+ * Returns a list of diagnostics based on the assertions provided.
+ *
+ * @param typeChecker - The TypeScript type checker.
+ * @param assertions - Assertion map with the key being the assertion, and the value the list of all those assertion nodes.
+ * @returns List of diagnostics.
+ */
+export const handle = (typeChecker: TypeChecker, assertions: Map<Assertion, Set<CallExpression>>): Diagnostic[] => {
+	const diagnostics: Diagnostic[] = [];
+
+	for (const [assertion, nodes] of assertions) {
+		const handler = assertionHandlers.get(assertion);
+
+		if (!handler) {
+			// Ignore these assertions as no handler is found
+			continue;
+		}
+
+		const handlers = Array.isArray(handler) ? handler : [handler];
+
+		// Iterate over the handlers and invoke them
+		for (const fn of handlers) {
+			diagnostics.push(...fn(typeChecker, nodes));
+		}
+	}
+
+	return diagnostics;
+};

source/lib/compiler.ts

@@ -3,16 +3,11 @@ import {
 	flattenDiagnosticMessageText,
 	createProgram,
 	Diagnostic as TSDiagnostic,
-	Program,
-	SourceFile,
-	Node,
-	forEachChild,
-	isCallExpression,
-	Identifier,
-	TypeChecker,
-	CallExpression
+	SourceFile
 } from '../../libraries/typescript';
+import {extractAssertions, parseErrorAssertionToLocation} from './parser';
 import {Diagnostic, DiagnosticCode, Context, Location} from './interfaces';
+import {handle} from './assertions';
 
 // List of diagnostic codes that should be ignored in general
 const ignoredDiagnostics = new Set<number>([
@@ -30,121 +25,12 @@ const diagnosticCodesToIgnore = new Set<DiagnosticCode>([
 	DiagnosticCode.NoOverloadMatches
 ]);
 
-/**
- * Extract all assertions.
- *
- * @param program - TypeScript program.
- */
-const extractAssertions = (program: Program) => {
-	const typeAssertions = new Set<CallExpression>();
-	const errorAssertions = new Set<CallExpression>();
-
-	function walkNodes(node: Node) {
-		if (isCallExpression(node)) {
-			const text = (node.expression as Identifier).getText();
-
-			if (text === 'expectType') {
-				typeAssertions.add(node);
-			} else if (text === 'expectError') {
-				errorAssertions.add(node);
-			}
-		}
-
-		forEachChild(node, walkNodes);
-	}
-
-	for (const sourceFile of program.getSourceFiles()) {
-		walkNodes(sourceFile);
-	}
-
-	return {
-		typeAssertions,
-		errorAssertions
-	};
-};
-
-/**
- * Loop over all the `expectError` nodes and convert them to a range map.
- *
- * @param nodes - Set of `expectError` nodes.
- */
-const extractExpectErrorRanges = (nodes: Set<Node>) => {
-	const expectedErrors = new Map<Location, Pick<Diagnostic, 'fileName' | 'line' | 'column'>>();
-
-	// Iterate over the nodes and add the node range to the map
-	for (const node of nodes) {
-		const location = {
-			fileName: node.getSourceFile().fileName,
-			start: node.getStart(),
-			end: node.getEnd()
-		};
-
-		const pos = node
-			.getSourceFile()
-			.getLineAndCharacterOfPosition(node.getStart());
-
-		expectedErrors.set(location, {
-			fileName: location.fileName,
-			line: pos.line + 1,
-			column: pos.character
-		});
-	}
-
-	return expectedErrors;
-};
-
-/**
- * Assert the expected type from `expectType` calls with the provided type in the argument.
- * Returns a list of custom diagnostics.
- *
- * @param checker - The TypeScript type checker.
- * @param nodes - The `expectType` AST nodes.
- * @return List of custom diagnostics.
- */
-const assertTypes = (checker: TypeChecker, nodes: Set<CallExpression>): Diagnostic[] => {
-	const diagnostics: Diagnostic[] = [];
-
-	for (const node of nodes) {
-		if (!node.typeArguments) {
-			// Skip if the node does not have generics
-			continue;
-		}
-
-		// Retrieve the type to be expected. This is the type inside the generic.
-		const expectedType = checker.getTypeFromTypeNode(node.typeArguments[0]);
-		const argumentType = checker.getTypeAtLocation(node.arguments[0]);
-
-		if (!checker.isAssignableTo(argumentType, expectedType)) {
-			// The argument type is not assignable to the expected type. TypeScript will catch this for us.
-			continue;
-		}
-
-		if (!checker.isAssignableTo(expectedType, argumentType)) { // tslint:disable-line:early-exit
-			/**
-			 * At this point, the expected type is not assignable to the argument type, but the argument type is
-			 * assignable to the expected type. This means our type is too wide.
-			 */
-			const position = node.getSourceFile().getLineAndCharacterOfPosition(node.getStart());
-
-			diagnostics.push({
-				fileName: node.getSourceFile().fileName,
-				message: `Parameter type \`${checker.typeToString(expectedType)}\` is declared too wide for argument type \`${checker.typeToString(argumentType)}\`.`,
-				severity: 'error',
-				line: position.line + 1,
-				column: position.character,
-			});
-		}
-	}
-
-	return diagnostics;
-};
-
 /**
  * Check if the provided diagnostic should be ignored.
  *
  * @param diagnostic - The diagnostic to validate.
  * @param expectedErrors - Map of the expected errors.
- * @return Boolean indicating if the diagnostic should be ignored or not.
+ * @returns Boolean indicating if the diagnostic should be ignored or not.
  */
 const ignoreDiagnostic = (diagnostic: TSDiagnostic, expectedErrors: Map<Location, any>): boolean => {
 	if (ignoredDiagnostics.has(diagnostic.code)) {
@@ -180,28 +66,28 @@ const ignoreDiagnostic = (diagnostic: TSDiagnostic, expectedErrors: Map<Location
 export const getDiagnostics = (context: Context): Diagnostic[] => {
 	const fileNames = context.testFiles.map(fileName => path.join(context.cwd, fileName));
 
-	const result: Diagnostic[] = [];
+	const diagnostics: Diagnostic[] = [];
 
 	const program = createProgram(fileNames, context.config.compilerOptions);
 
-	const diagnostics = program
+	const tsDiagnostics = program
 		.getSemanticDiagnostics()
 		.concat(program.getSyntacticDiagnostics());
 
-	const {typeAssertions, errorAssertions} = extractAssertions(program);
+	const assertions = extractAssertions(program);
 
-	const expectedErrors = extractExpectErrorRanges(errorAssertions);
+	diagnostics.push(...handle(program.getTypeChecker(), assertions));
 
-	result.push(...assertTypes(program.getTypeChecker(), typeAssertions));
+	const expectedErrors = parseErrorAssertionToLocation(assertions);
 
-	for (const diagnostic of diagnostics) {
+	for (const diagnostic of tsDiagnostics) {
 		if (!diagnostic.file || ignoreDiagnostic(diagnostic, expectedErrors)) {
 			continue;
 		}
 
 		const position = diagnostic.file.getLineAndCharacterOfPosition(diagnostic.start as number);
 
-		result.push({
+		diagnostics.push({
 			fileName: diagnostic.file.fileName,
 			message: flattenDiagnosticMessageText(diagnostic.messageText, '\n'),
 			severity: 'error',
@@ -211,12 +97,12 @@ export const getDiagnostics = (context: Context): Diagnostic[] => {
 	}
 
 	for (const [, diagnostic] of expectedErrors) {
-		result.push({
+		diagnostics.push({
 			...diagnostic,
 			message: 'Expected an error, but found none.',
 			severity: 'error'
 		});
 	}
 
-	return result;
+	return diagnostics;
 };

source/lib/parser.ts

@@ -0,0 +1,80 @@
+import {Program, Node, CallExpression, forEachChild, isCallExpression, Identifier} from '../../libraries/typescript';
+import {Assertion} from './assertions';
+import {Location, Diagnostic} from './interfaces';
+
+// TODO: Use Object.values() when targetting Node.js >= 8
+const assertionTypes = new Set<string>(Object.keys(Assertion).map(key => Assertion[key]));
+
+/**
+ * Extract all assertions.
+ *
+ * @param program - TypeScript program.
+ */
+export const extractAssertions = (program: Program): Map<Assertion, Set<CallExpression>> => {
+	const assertions = new Map<Assertion, Set<CallExpression>>();
+
+	/**
+	 * Recursively loop over all the nodes and extract all the assertions out of the source files.
+	 */
+	function walkNodes(node: Node) {
+		if (isCallExpression(node)) {
+			const text = (node.expression as Identifier).getText();
+
+			// Check if the call type is a valid assertion
+			if (assertionTypes.has(text)) {
+				const assertion = text as Assertion;
+
+				const nodes = assertions.get(assertion) || new Set<CallExpression>();
+
+				nodes.add(node);
+
+				assertions.set(assertion, nodes);
+			}
+		}
+
+		forEachChild(node, walkNodes);
+	}
+
+	for (const sourceFile of program.getSourceFiles()) {
+		walkNodes(sourceFile);
+	}
+
+	return assertions;
+};
+
+/**
+ * Loop over all the error assertion nodes and convert them to a location map.
+ *
+ * @param assertions - Assertion map.
+ */
+export const parseErrorAssertionToLocation = (assertions: Map<Assertion, Set<CallExpression>>) => {
+	const nodes = assertions.get(Assertion.EXPECT_ERROR);
+
+	const expectedErrors = new Map<Location, Pick<Diagnostic, 'fileName' | 'line' | 'column'>>();
+
+	if (!nodes) {
+		// Bail out if we don't have any error nodes
+		return expectedErrors;
+	}
+
+	// Iterate over the nodes and add the node range to the map
+	for (const node of nodes) {
+		const location = {
+			fileName: node.getSourceFile().fileName,
+			start: node.getStart(),
+			end: node.getEnd()
+		};
+
+		const pos = node
+			.getSourceFile()
+			.getLineAndCharacterOfPosition(node.getStart());
+
+		expectedErrors.set(location, {
+			fileName: location.fileName,
+			line: pos.line + 1,
+			column: pos.character
+		});
+	}
+
+	return expectedErrors;
+};