Merge pull request #1 from code4rena-dev/nlf/discriminated-unions

feat: make responses discriminated unions for stronger type safety while not throwing for invalid status codes

Author: nlf

lib/build/index.ts

@@ -1,12 +1,15 @@
 import type { RequestRoute, Server } from "@hapi/hapi";
 import Joi from "joi";
 import type ts from "typescript";
-import { addSyntheticLeadingComment, factory as f } from "typescript";
+import { addSyntheticLeadingComment } from "typescript";
 
 import type { ExtendedObjectSchema, ExtendedSchema } from "./joi.ts";
 import {
+  importDeclaration,
   typeAliasDeclaration,
   typeLiteralNode,
+  typeReferenceNode,
+  unionTypeNode,
 } from "./nodes.ts";
 import {
   generateType,
@@ -51,9 +54,11 @@ export function getTypeName(routeName: string, suffix: string): string {
 }
 
 export function generateClientType(server: Server) {
-  const routeMap: Record<string, Record<string, Record<string, { name: string; required: boolean }>>> = {};
+  const routeList: ts.TypeNode[] = [];
   const statements: ts.Statement[] = [];
 
+  statements.push(importDeclaration(["Route", "StatusCode"], "@code4rena/typed-client", true));
+
   for (const route of server.table()) {
     const routeName = getRouteName(route);
     const routeOptions: Record<string, { name: string; required: boolean }> = {};
@@ -108,40 +113,62 @@ export function generateClientType(server: Server) {
       statements.push(typeAliasDeclaration(payloadTypeName, generateType(route.settings.validate.payload as ExtendedObjectSchema), true));
     }
 
+    // the result type is the type of the actual response from the api
+    // const resultTypeName = getTypeName(routeName, "result");
+    // the response type is the result wrapped in the higher level object including url, status, etc.
     const responseTypeName = getTypeName(routeName, "response");
-    const responseValidator = Object.entries(route.settings.response?.status ?? {})
-      .filter(([code]) => +code >= 200 && +code < 300)
-      .map(([_, schema]) => schema)[0] ?? Joi.any();
+    const matchedCodes: string[] = [];
+    const responseTypeList: string[] = [];
+    // first, we iterate our response validators and get everything that has a specific response code
+    for (const [code, schema] of Object.entries(route.settings.response?.status ?? {})) {
+      matchedCodes.push(code);
+
+      const responseCodeTypeName = getTypeName(responseTypeName, code);
+      const responseNode = typeLiteralNode({
+        status: { name: `${code}`, required: true },
+        ok: { name: (+code >= 200 && +code < 300) ? "true" : "false", required: true },
+        headers: { name: "Headers", required: true },
+        url: { name: "string", required: true },
+        data: { node: generateType(schema as ExtendedObjectSchema), required: isRequired(schema as ExtendedSchema) },
+      });
+      statements.push(typeAliasDeclaration(responseCodeTypeName, responseNode, true));
+      responseTypeList.push(responseCodeTypeName);
+    }
+
+    // now we insert a final response type where the status code is Exclude<StatusCodes, matchedCodes>
+    const unknownResponseName = getTypeName(responseTypeName, "Unknown");
+    const unknownResponseNode = typeLiteralNode({
+      status: {
+        name: matchedCodes.length
+          // HACK: this could probably build a real node with a real union passed to a real generic
+          ? `Exclude<StatusCode, ${matchedCodes.join(" | ")}>`
+          : "StatusCode",
+        required: true,
+      },
+      ok: { name: "boolean", required: true },
+      headers: { name: "Headers", required: true },
+      url: { name: "string", required: true },
+      data: { name: "unknown", required: false },
+    });
+    statements.push(typeAliasDeclaration(unknownResponseName, unknownResponseNode, true));
+    responseTypeList.push(unknownResponseName);
+
+    const responseUnionType = unionTypeNode(responseTypeList.map((responseType) => typeReferenceNode(responseType)));
+    statements.push(typeAliasDeclaration(responseTypeName, responseUnionType, true));
+
     routeOptions.response = {
       name: responseTypeName,
       required: true,
     };
-    statements.push(typeAliasDeclaration(responseTypeName, generateType(responseValidator as ExtendedObjectSchema), true));
 
-    routeMap[route.method.toUpperCase()] ??= {};
-    routeMap[route.method.toUpperCase()][route.path] = routeOptions;
+    const routeType = typeReferenceNode("Route", route.method.toUpperCase(), route.path, typeLiteralNode(routeOptions));
+    routeList.push(routeType);
   }
 
-  const clientTypeNode = f.createTypeLiteralNode(Object.entries(routeMap).map(([method, route]) => {
-    return f.createPropertySignature(
-      [],
-      f.createStringLiteral(method),
-      undefined,
-      f.createTypeLiteralNode(Object.entries(route).map(([path, options]) => {
-        return f.createPropertySignature(
-          [],
-          f.createStringLiteral(path),
-          undefined,
-          typeLiteralNode(options),
-        );
-      })));
-  }));
-
-  statements.push(typeAliasDeclaration("RouteMap", clientTypeNode, true));
-
-  addSyntheticLeadingComment(statements[0], SyntaxKind.MultiLineCommentTrivia, " eslint-disable @stylistic/indent ", true);
-  addSyntheticLeadingComment(statements[0], SyntaxKind.MultiLineCommentTrivia, " eslint-disable @stylistic/quote-props ", true);
-  addSyntheticLeadingComment(statements[0], SyntaxKind.MultiLineCommentTrivia, " eslint-disable @stylistic/comma-dangle ", true);
+  statements.push(typeAliasDeclaration("Routes", unionTypeNode(routeList), true));
+
+  // throw a comment at the front to disable eslint for the file since it may not align with formatting
+  addSyntheticLeadingComment(statements[0], SyntaxKind.MultiLineCommentTrivia, " eslint-disable ", true);
 
   const sourceFile = factory.createSourceFile(statements, factory.createToken(SyntaxKind.EndOfFileToken), NodeFlags.None);
   const printer = createPrinter();

lib/build/nodes.ts

@@ -9,6 +9,7 @@ import type {
   QuestionToken,
   Statement,
   TypeNode,
+  TypeParameterDeclaration,
 } from "typescript";
 import {
   EmitHint,
@@ -54,15 +55,13 @@ export function dedupeNodes(nodes: TypeNode[]): TypeNode[] {
   });
 }
 
-export function importDeclaration(name: string, from: string, typeOnly?: boolean) {
+export function importDeclaration(name: string | string[], from: string, typeOnly?: boolean) {
   return f.createImportDeclaration(
     undefined,
     f.createImportClause(
       typeOnly ?? false,
       undefined,
-      f.createNamedImports([
-        f.createImportSpecifier(false, undefined, f.createIdentifier(name)),
-      ]),
+      f.createNamedImports([...name].map((name) => f.createImportSpecifier(false, undefined, f.createIdentifier(name)))),
     ),
     f.createStringLiteral(from),
   );
@@ -81,8 +80,12 @@ export function exportDeclaration(name: string, from: string) {
   );
 }
 
-export function typeAliasDeclaration(name: string, node: TypeNode, exported: boolean = false) {
-  return f.createTypeAliasDeclaration(exported ? [exportToken()] : [], name, undefined, node);
+export function typeAliasDeclaration(name: string, node: TypeNode, exported: boolean = false, types: TypeParameterDeclaration[] = []) {
+  return f.createTypeAliasDeclaration(exported ? [exportToken()] : [], name, types, node);
+}
+
+export function typeParameterDeclaration(name: string, constraint: TypeNode = undefined, def: TypeNode = undefined) {
+  return f.createTypeParameterDeclaration(undefined, name, constraint, def);
 }
 
 export function questionToken(schema?: Schema): QuestionToken | undefined {
@@ -149,21 +152,29 @@ export function unknownTypeNode() {
   return f.createKeywordTypeNode(SyntaxKind.UnknownKeyword);
 }
 
-export function typeLiteralNode(props: Record<string, { name: string; required: boolean }>) {
+type typeLiteralNodeProps = Record<string, {
+  required: boolean;
+  types?: TypeNode[];
+} & ({ name: string; node?: never } | { name?: never; node: TypeNode })>;
+export function typeLiteralNode(props: typeLiteralNodeProps) {
   return f.createTypeLiteralNode(Object.entries(props).map(([key, prop]) => {
     return f.createPropertySignature(
       [],
-      f.createStringLiteral(key),
+      f.createIdentifier(key),
       prop.required ? undefined : questionToken(),
-      f.createTypeReferenceNode(f.createIdentifier(prop.name)),
+      typeof prop.name === "string"
+        ? f.createTypeReferenceNode(f.createIdentifier(prop.name), prop.types)
+        : prop.node,
     );
   }));
 }
 
-export function typeReferenceNode(name: string, ...typeArguments: string[]) {
+export function typeReferenceNode(name: string, ...typeArguments: (string | TypeNode)[]) {
   return f.createTypeReferenceNode(
     f.createIdentifier(name),
-    typeArguments.map((typeArgument) => f.createTypeReferenceNode(f.createIdentifier(typeArgument))),
+    typeArguments.map((typeArgument) => typeof typeArgument === "string"
+      ? f.createLiteralTypeNode(f.createStringLiteral(typeArgument))
+      : typeArgument),
   );
 }
 
@@ -173,13 +184,13 @@ export function objectTypeNode(props: Property[] | Map<string, Schema>, options:
   const mappedProps = Array.isArray(props)
     ? props.map((prop) => f.createPropertySignature(
         [],
-        f.createStringLiteral(prop.key),
+        f.createIdentifier(prop.key),
         questionToken(prop.schema),
         generateType(prop.schema, _options)),
       )
     : Array.from(props.entries()).map(([key, schema]) => f.createPropertySignature(
         [],
-        f.createStringLiteral(key),
+        f.createIdentifier(key),
         questionToken(schema),
         generateType(schema, _options),
       ));