feat: assertion api detection

src/api/tsd/index.ts

@@ -0,0 +1,5 @@
+export const VITE_PLUGIN_VITEST_TYPESCRIPT_ASSERT = 'tsd';
+
+export function expectType() {
+  return;
+}

src/api/tssert/index.ts

@@ -0,0 +1,5 @@
+export const VITE_PLUGIN_VITEST_TYPESCRIPT_ASSERT = 'tssert';
+
+export function expectType() {
+  return;
+}

src/error.ts

@@ -1,11 +1,16 @@
 export enum ErrorCode {
+  UNEXPECTED_ERROR,
   TSCONFIG_FILE_NOT_FOUND,
   TSCONFIG_FILE_NOT_READABLE,
+  MULTIPLE_ASSERTION_API_IN_SAME_FILE_NOT_ALLOWED,
 }
 
 export const errorMessages: Record<ErrorCode, string> = {
+  [ErrorCode.UNEXPECTED_ERROR]: 'Unexpected error: {message}',
   [ErrorCode.TSCONFIG_FILE_NOT_FOUND]: 'TS config file not found. File name: {configName} - Search path: {searchPath}',
   [ErrorCode.TSCONFIG_FILE_NOT_READABLE]: 'TS config file not readable or empty. File path: {filePath}',
+  [ErrorCode.MULTIPLE_ASSERTION_API_IN_SAME_FILE_NOT_ALLOWED]:
+    'The use of multiple type assertion APIs in the same file is not allowed. The APIs detected are [{apiNames}], please review your test.',
 };
 
 export function errorMessage(code: ErrorCode, data?: Record<string, unknown>) {

src/plugin/error.ts

@@ -19,11 +19,11 @@ export function createErrorString({
   err.stacks = [{
     file: "${file}",
     line: ${line + 1},
-    column: ${column},
+    column: ${column + 1},
     sourcePos: {
       source: "${file}",
       line: ${line + 1},
-      column: ${column}
+      column: ${column + 1}
     }
   }];
   throw err;

src/plugin/transform.ts

@@ -1,10 +1,11 @@
-import type ts from 'byots';
+import ts from 'byots';
 import type { Reporter } from '.';
 import MagicString from 'magic-string';
-import { reportDiagnostics } from './util';
 import type { TransformResult } from 'vite';
-import { program } from '../typescript/program';
+import { ErrorCode, errorMessage } from '../error';
+import { createProgram } from '../typescript/program';
 import type { TypeScriptConfigOptions } from '../typescript/config';
+import { type APIName, reportDiagnostics, tryToGetAPINAme } from './util';
 
 export interface TransformSettings {
   code: string;
@@ -17,15 +18,50 @@ export interface TransformSettings {
 }
 
 export function transform({ code, fileName, report, typescript }: TransformSettings): TransformResult {
-  const { diagnostics } = program({ config: typescript.config, fileName });
+  const { diagnostics, checker, sourceFile, program } = createProgram({ config: typescript.config, fileName });
   const newCode = new MagicString(code);
 
   if (report.includes('type-error')) {
     reportDiagnostics(diagnostics, newCode, fileName);
   }
 
+  if (report.includes('type-assertion')) {
+    typeAssertion(fileName, sourceFile, program, checker);
+  }
+
   return {
     code: newCode.toString(),
     map: newCode.generateMap({ hires: true }),
   };
 }
+
+function typeAssertion(fileName: string, sourceFile: ts.SourceFile, program: ts.Program, checker: ts.TypeChecker) {
+  let apiName: APIName | undefined = undefined;
+
+  function visit(node: ts.Node) {
+    if (ts.isImportDeclaration(node)) {
+      const moduleName = node.moduleSpecifier.getText().slice(1, -1);
+      const moduleAPIName = tryToGetAPINAme(moduleName, fileName, program, checker);
+
+      if (apiName && moduleAPIName && apiName !== moduleAPIName) {
+        throw errorMessage(ErrorCode.MULTIPLE_ASSERTION_API_IN_SAME_FILE_NOT_ALLOWED, {
+          apiNames: `${apiName}, ${moduleAPIName}`,
+        });
+      }
+
+      apiName = moduleAPIName;
+
+      // eslint-disable-next-line no-console
+      console.log({ apiName });
+    }
+
+    if (ts.isCallExpression(node)) {
+      // eslint-disable-next-line no-console
+      console.log('>', node.getText());
+    }
+
+    node.forEachChild(visit);
+  }
+
+  visit(sourceFile);
+}

src/plugin/util.ts

@@ -1,8 +1,13 @@
 import ts from 'byots';
 import type MagicString from 'magic-string';
 import { createErrorString } from './error';
-import { newLine } from '../typescript/util';
 import { testWrapperIdentifiers } from './identifiers';
+import { getResolvedModuleExports, newLine } from '../typescript/util';
+
+export type APIName = 'tsd' | 'tssert';
+
+export const API_MODULE_KEY = 'VITE_PLUGIN_VITEST_TYPESCRIPT_ASSERT';
+export const API_NAMES: readonly APIName[] = ['tsd', 'tssert'] as const;
 
 export function searchTestWrapperFromPosition(file: ts.SourceFile, position: number) {
   const token = ts.getTokenAtPosition(file, position);
@@ -51,3 +56,23 @@ export function reportDiagnostics(diagnostics: readonly ts.Diagnostic[], newCode
     }
   }
 }
+
+// Probably the most disgusting function name I've ever written in my life. Pushing the limits!
+export function tryToGetAPINAme(moduleName: string, fileName: string, program: ts.Program, checker: ts.TypeChecker) {
+  const resolvedExports = getResolvedModuleExports(program, moduleName, fileName, checker);
+
+  if (!resolvedExports) {
+    return;
+  }
+
+  const apiKey = resolvedExports.find((e) => e.escapedName === API_MODULE_KEY);
+
+  if (!apiKey?.declarations?.[0]) {
+    return;
+  }
+
+  const apiType = checker.getTypeOfSymbolAtLocation(apiKey, apiKey.declarations[0]);
+  const apiName = checker.typeToString(apiType).slice(1, -1) as APIName;
+
+  return API_NAMES.includes(apiName) ? apiName : undefined;
+}

src/typescript/program.ts

@@ -1,10 +1,18 @@
 import ts from 'byots';
+import { ErrorCode, errorMessage } from '../error';
 
-export function program({ fileName, config }: { fileName: string; config: ts.ParsedCommandLine }) {
+export function createProgram({ fileName, config }: { fileName: string; config: ts.ParsedCommandLine }) {
   const compilerOptions = { ...ts.getDefaultCompilerOptions(), ...config.options };
   const program = ts.createProgram([fileName], compilerOptions);
   const diagnostics = ts.getPreEmitDiagnostics(program);
+  const sourceFile = program.getSourceFile(fileName);
   const checker = program.getTypeChecker();
 
-  return { program, diagnostics, checker };
+  if (!sourceFile) {
+    throw errorMessage(ErrorCode.UNEXPECTED_ERROR, {
+      message: `The source file is unreachable. File path: ${fileName})`,
+    });
+  }
+
+  return { compilerOptions, program, diagnostics, sourceFile, checker };
 }

src/typescript/util.ts

@@ -18,3 +18,30 @@ export function getMiddle(node: ts.Node) {
   const diff = node.getEnd() - node.getStart();
   return node.getStart() + Math.round(diff / 2);
 }
+
+export function getResolvedModuleExports(
+  program: ts.Program,
+  moduleName: string,
+  fileName: string,
+  checker: ts.TypeChecker,
+) {
+  const { resolvedModule } = program.getResolvedModuleWithFailedLookupLocationsFromCache(moduleName, fileName) ?? {};
+
+  if (!resolvedModule) {
+    return;
+  }
+
+  const moduleSourceFile = program.getSourceFile(resolvedModule.resolvedFileName);
+
+  if (!moduleSourceFile) {
+    return;
+  }
+
+  const moduleSourceFileSymbol = checker.getSymbolAtLocation(moduleSourceFile);
+
+  if (!moduleSourceFileSymbol) {
+    return;
+  }
+
+  return checker.getExportsOfModule(moduleSourceFileSymbol);
+}

test/index.test.ts

@@ -1,12 +1,17 @@
 import { test, expect, describe } from 'vitest';
 
+import { expectType } from '../src/api/tsd';
+// import { expectType } from '../src/api/tssert';
+
 // eslint-disable-next-line @typescript-eslint/no-unused-vars
 const prout1 = 'plop';
 
 test('test-1', () => {
   expect('Hello World').toBe(42);
 
   // const prout2 = 'plop';
+
+  expectType();
 });
 
 describe('describe-1', () => {