feat(parser): simplify ParseResult and parse inference (#3568)

Author: am29d

Diff for: packages/parser/src/parser.ts

@@ -1,6 +1,7 @@
 import type { ZodSchema, z } from 'zod';
 import { ParseError } from './errors.js';
-import type { Envelope, ParsedResult } from './types/index.js';
+import type { Envelope } from './types/index.js';
+import type { ParseFunction } from './types/parser.js';
 
 /**
  * Parse the data using the provided schema, envelope and safeParse flag
@@ -27,12 +28,12 @@ import type { Envelope, ParsedResult } from './types/index.js';
  * @param schema the schema to use
  * @param safeParse whether to use safeParse or not, if true it will return a ParsedResult with the original event if the parsing fails
  */
-const parse = <T extends ZodSchema, E extends Envelope>(
+const parse: ParseFunction = <T extends ZodSchema, E extends Envelope>(
   data: z.infer<T>,
   envelope: E | undefined,
   schema: T,
   safeParse?: boolean
-): ParsedResult | z.infer<T> => {
+) => {
   if (envelope && safeParse) {
     return envelope.safeParse(data, schema);
   }
@@ -56,10 +57,7 @@ const parse = <T extends ZodSchema, E extends Envelope>(
  * @param data the data to parse
  * @param schema the zod schema to use
  */
-const safeParseSchema = <T extends ZodSchema>(
-  data: z.infer<T>,
-  schema: T
-): ParsedResult => {
+const safeParseSchema = <T extends ZodSchema>(data: z.infer<T>, schema: T) => {
   const result = schema.safeParse(data);
 
   return result.success

Diff for: packages/parser/src/types/index.ts

@@ -1,42 +1,43 @@
 export type {
-  ParserOptions,
   ParsedResult,
-  ParsedResultSuccess,
   ParsedResultError,
+  ParsedResultSuccess,
+  ParseFunction,
+  ParserOptions,
 } from '../types/parser.js';
 export type { Envelope } from './envelope.js';
 
 export type {
   ALBEvent,
-  APIGatewayProxyEvent,
   ALBMultiValueHeadersEvent,
+  APIGatewayProxyEvent,
   APIGatewayProxyEventV2,
-  APIGatewayRequestContextV2,
   APIGatewayRequestAuthorizerV2,
-  AppSyncResolverEvent,
+  APIGatewayRequestContextV2,
   AppSyncBatchResolverEvent,
-  S3Event,
-  S3EventNotificationEventBridge,
-  S3SqsEventNotification,
-  SnsEvent,
-  SqsEvent,
-  DynamoDBStreamEvent,
-  DynamoDBStreamToKinesisRecordEvent,
-  CloudWatchLogsEvent,
+  AppSyncResolverEvent,
   CloudFormationCustomResourceCreateEvent,
   CloudFormationCustomResourceDeleteEvent,
   CloudFormationCustomResourceUpdateEvent,
+  CloudWatchLogsEvent,
+  DynamoDBStreamEvent,
+  DynamoDBStreamToKinesisRecordEvent,
   EventBridgeEvent,
-  KafkaSelfManagedEvent,
   KafkaMskEvent,
+  KafkaSelfManagedEvent,
   KinesisDataStreamEvent,
   KinesisDynamoDBStreamEvent,
   KinesisFireHoseEvent,
   KinesisFireHoseSqsEvent,
   LambdaFunctionUrlEvent,
+  S3Event,
+  S3EventNotificationEventBridge,
+  S3ObjectLambdaEvent,
+  S3SqsEventNotification,
   SesEvent,
+  SnsEvent,
   SnsSqsNotification,
-  S3ObjectLambdaEvent,
+  SqsEvent,
   VpcLatticeEvent,
   VpcLatticeEventV2,
 } from './schema.js';

Diff for: packages/parser/src/types/parser.ts

@@ -34,7 +34,7 @@ type ParsedResultError<Input> = {
 /**
  * The result of parsing an event using the safeParse, can either be a success or an error
  */
-type ParsedResult<Input = unknown, Output = unknown> =
+type ParsedResult<Input = unknown, Output = Input> =
   | ParsedResultSuccess<Output>
   | ParsedResultError<Input>;
 
@@ -54,7 +54,7 @@ type ZodInferredSafeParseResult<
   TSchema extends ZodSchema,
   TEnvelope extends Envelope,
 > = undefined extends TEnvelope
-  ? ParsedResult<unknown, z.infer<TSchema>>
+  ? ParsedResult<z.infer<TSchema>, z.infer<TSchema>>
   : TEnvelope extends ArrayEnvelope
     ? ParsedResult<unknown, z.infer<TSchema>[]>
     : ParsedResult<unknown, z.infer<TSchema>>;
@@ -70,10 +70,63 @@ type ParserOutput<
   ? ZodInferredSafeParseResult<TSchema, TEnvelope>
   : ZodInferredResult<TSchema, TEnvelope>;
 
+/**
+ * The parser function that can parse the data using the provided schema and envelope
+ * we use function overloads to provide the correct return type based on the provided envelope
+ **/
+type ParseFunction = {
+  // No envelope cases
+  <T extends ZodSchema>(
+    data: z.infer<T>,
+    envelope: undefined,
+    schema: T,
+    safeParse?: false
+  ): z.infer<T>;
+
+  <T extends ZodSchema>(
+    data: z.infer<T>,
+    envelope: undefined,
+    schema: T,
+    safeParse: true
+  ): ParsedResult<z.infer<T>>;
+
+  // Generic envelope case
+  <T extends ZodSchema, E extends Envelope>(
+    data: unknown,
+    envelope: E,
+    schema: T,
+    safeParse?: false
+  ): E extends ArrayEnvelope ? z.infer<T>[] : z.infer<T>;
+
+  <T extends ZodSchema, E extends Envelope>(
+    data: unknown,
+    envelope: E,
+    schema: T,
+    safeParse: true
+  ): E extends ArrayEnvelope
+    ? ParsedResult<unknown, z.infer<T>[]>
+    : ParsedResult<unknown, z.infer<T>>;
+
+  // Generic envelope case with safeParse
+  <T extends ZodSchema, E extends Envelope, S extends boolean = false>(
+    data: unknown,
+    envelope: E,
+    schema: T,
+    safeParse?: S
+  ): S extends true
+    ? E extends ArrayEnvelope
+      ? ParsedResult<unknown, z.infer<T>[]>
+      : ParsedResult<unknown, z.infer<T>>
+    : E extends ArrayEnvelope
+      ? z.infer<T>[]
+      : z.infer<T>;
+};
+
 export type {
-  ParserOptions,
+  ParseFunction,
   ParsedResult,
   ParsedResultError,
   ParsedResultSuccess,
+  ParserOptions,
   ParserOutput,
 };

Diff for: packages/parser/tests/types/envelopes.test.ts

@@ -1,4 +1,4 @@
-import { describe, expect, it } from 'vitest';
+import { describe, expect, expectTypeOf, it } from 'vitest';
 import { z } from 'zod';
 import {
   ApiGatewayEnvelope,
@@ -69,10 +69,6 @@ describe('Types ', () => {
     expect(Array.isArray(result)).toBe(true);
     expect(result).toEqual([{ name: 'John', age: 30 }]);
 
-    // Type assertion to ensure it's specifically User[]
-    type AssertIsUserArray<T> = T extends z.infer<typeof userSchema>[]
-      ? true
-      : false;
-    type Test = AssertIsUserArray<Result>;
+    expectTypeOf(result).toEqualTypeOf<z.infer<typeof userSchema>[]>();
   });
 });

Diff for: packages/parser/tests/types/parser.test.ts

@@ -0,0 +1,88 @@
+import { describe } from 'node:test';
+import { expect, expectTypeOf, it } from 'vitest';
+import { z } from 'zod';
+import { EventBridgeEnvelope, SqsEnvelope } from '../../src/envelopes/index.js';
+import { parse } from '../../src/parser.js';
+import type { EventBridgeEvent, SqsEvent } from '../../src/types/schema.js';
+import { getTestEvent } from '../unit/helpers/utils.js';
+
+describe('Parser types', () => {
+  const userSchema = z.object({
+    name: z.string(),
+    age: z.number(),
+  });
+  type User = z.infer<typeof userSchema>;
+  const input = { name: 'John', age: 30 };
+
+  const eventBridgeBaseEvent = getTestEvent<EventBridgeEvent>({
+    eventsPath: 'eventbridge',
+    filename: 'base',
+  });
+
+  const sqsBaseEvent = getTestEvent<SqsEvent>({
+    eventsPath: 'sqs',
+    filename: 'base',
+  });
+  it('infers return type for schema and safeParse', () => {
+    // Act
+    const result = parse(input, undefined, userSchema, true);
+
+    // Assert
+    if (result.success) {
+      expectTypeOf(result.data).toEqualTypeOf<User>();
+    } else {
+      expectTypeOf(result.originalEvent).toEqualTypeOf<User | undefined>();
+    }
+  });
+
+  it('infers return type for schema', () => {
+    // Act
+    const result = parse(input, undefined, userSchema);
+
+    // Assert
+    expectTypeOf(result).toEqualTypeOf<User>();
+  });
+
+  it('infers return type for schema and envelope', () => {
+    // Prepare
+    const event = structuredClone(eventBridgeBaseEvent);
+    event.detail = input;
+
+    // Act
+    const result = parse(event, EventBridgeEnvelope, userSchema);
+
+    // Assert
+    expectTypeOf(result).toEqualTypeOf<User>();
+  });
+
+  it('infert return type for schema, object envelope and safeParse', () => {
+    // Prepare
+    const event = structuredClone(eventBridgeBaseEvent);
+    event.detail = input;
+
+    // Act
+    const result = parse(event, EventBridgeEnvelope, userSchema, true);
+
+    // Assert
+    if (result.success) {
+      expectTypeOf(result.data).toEqualTypeOf<User>();
+      expect(result.data).toEqual(input);
+    } else {
+      throw new Error('Parsing failed');
+    }
+  });
+
+  it('infers return type for schema, array envelope and safeParse', () => {
+    // Prepare
+    const event = structuredClone(sqsBaseEvent);
+    event.Records[0].body = JSON.stringify(input);
+
+    // Act
+    const result = parse(input, SqsEnvelope, userSchema, true);
+
+    // Assert
+    if (result.success) {
+      expectTypeOf(result.data).toEqualTypeOf<User[]>();
+    }
+  });
+});