add type inference for middy middleware

Author: am29d

packages/parser/src/middleware/parser.ts

@@ -1,8 +1,9 @@
 import { type MiddyLikeRequest } from '@aws-lambda-powertools/commons/types';
 import { type MiddlewareObj } from '@middy/core';
-import { type ZodSchema } from 'zod';
-import { type ParserOptions } from '../types/parser.js';
+import { ZodType } from 'zod';
+import type { ParserOptions, ParseOutput } from '../types/parser.js';
 import { parse } from '../parser.js';
+import type { Envelope } from '../types/envelope.js';
 
 /**
  * A middiy middleware to parse your event.
@@ -32,9 +33,13 @@ import { parse } from '../parser.js';
  *
  * @param options
  */
-const parser = <S extends ZodSchema>(
-  options: ParserOptions<S>
-): MiddlewareObj => {
+const parser = <
+  TSchema extends ZodType,
+  TEnvelope extends Envelope = undefined,
+  TSafeParse extends boolean = false,
+>(
+  options: ParserOptions<TSchema, TEnvelope, TSafeParse>
+): MiddlewareObj<ParseOutput<TSchema, TEnvelope, TSafeParse>> => {
   const before = (request: MiddyLikeRequest): void => {
     const { schema, envelope, safeParse } = options;
 

packages/parser/src/parserDecorator.ts

@@ -1,8 +1,9 @@
 import type { HandlerMethodDecorator } from '@aws-lambda-powertools/commons/types';
 import type { Context, Handler } from 'aws-lambda';
-import { ZodSchema, z } from 'zod';
+import { type ZodSchema } from 'zod';
 import { parse } from './parser.js';
-import type { ParserOptions, ParsedResult } from './types/index.js';
+import type { ParserOptions, Envelope } from './types/index.js';
+import type { ParseOutput } from './types/parser.js';
 
 /**
  * A decorator to parse your event.
@@ -67,8 +68,12 @@ import type { ParserOptions, ParsedResult } from './types/index.js';
  *
  * @param options Configure the parser with the `schema`, `envelope` and whether to `safeParse` or not
  */
-export const parser = <S extends ZodSchema>(
-  options: ParserOptions<S>
+export const parser = <
+  TSchema extends ZodSchema,
+  TEnvelope extends Envelope = undefined,
+  TSafeParse extends boolean = false,
+>(
+  options: ParserOptions<TSchema, TEnvelope, TSafeParse>
 ): HandlerMethodDecorator => {
   return (_target, _propertyKey, descriptor) => {
     const original = descriptor.value!;
@@ -77,14 +82,11 @@ export const parser = <S extends ZodSchema>(
 
     descriptor.value = async function (
       this: Handler,
-      event: unknown,
+      event: ParseOutput<TSchema, TEnvelope, TSafeParse>,
       context: Context,
       callback
     ) {
-      const parsedEvent: ParsedResult<
-        typeof event,
-        z.infer<typeof schema>
-      > = parse(event, envelope, schema, safeParse);
+      const parsedEvent = parse(event, envelope, schema, safeParse);
 
       return original.call(this, parsedEvent, context, callback);
     };

packages/parser/src/types/envelope.ts

@@ -14,8 +14,14 @@ import type {
   VpcLatticeEnvelope,
   VpcLatticeV2Envelope,
 } from '../envelopes/index.js';
+import { z, type ZodSchema } from 'zod';
 
-export type Envelope =
+type DynamoDBStreamEnvelopeResponse<Schema extends ZodSchema> = {
+  NewImage: z.infer<Schema>;
+  OldImage: z.infer<Schema>;
+};
+
+type Envelope =
   | typeof ApiGatewayEnvelope
   | typeof ApiGatewayV2Envelope
   | typeof CloudWatchEnvelope
@@ -29,4 +35,23 @@ export type Envelope =
   | typeof SnsSqsEnvelope
   | typeof SqsEnvelope
   | typeof VpcLatticeEnvelope
-  | typeof VpcLatticeV2Envelope;
+  | typeof VpcLatticeV2Envelope
+  | undefined;
+
+/**
+ * Envelopes that return an array, needed to narrow down the return type of the parser
+ */
+type EnvelopeArrayReturnType =
+  | typeof CloudWatchEnvelope
+  | typeof DynamoDBStreamEnvelope
+  | typeof KafkaEnvelope
+  | typeof KinesisEnvelope
+  | typeof KinesisFirehoseEnvelope
+  | typeof SnsEnvelope
+  | typeof SqsEnvelope;
+
+export type {
+  Envelope,
+  DynamoDBStreamEnvelopeResponse,
+  EnvelopeArrayReturnType,
+};

packages/parser/src/types/parser.ts

@@ -1,13 +1,17 @@
-import type { ZodSchema, ZodError } from 'zod';
-import type { Envelope } from './envelope.js';
+import { type ZodSchema, type ZodError, z } from 'zod';
+import type { Envelope, EnvelopeArrayReturnType } from './envelope.js';
 
 /**
  * Options for the parser used in middy middleware and decorator
  */
-type ParserOptions<S extends ZodSchema> = {
-  schema: S;
-  envelope?: Envelope;
-  safeParse?: boolean;
+type ParserOptions<
+  TSchema extends ZodSchema,
+  TEnvelope extends Envelope,
+  TSafeParse extends boolean,
+> = {
+  schema: TSchema;
+  envelope?: TEnvelope;
+  safeParse?: TSafeParse;
 };
 
 /**
@@ -34,9 +38,46 @@ type ParsedResult<Input = unknown, Output = unknown> =
   | ParsedResultSuccess<Output>
   | ParsedResultError<Input>;
 
+/**
+ * The inferred result of the schema, can be either an array or a single object depending on the envelope
+ */
+type ZodInferredResult<
+  TSchema extends ZodSchema,
+  TEnvelope extends Envelope,
+> = undefined extends TEnvelope
+  ? z.infer<TSchema>
+  : TEnvelope extends EnvelopeArrayReturnType
+    ? z.infer<TSchema>[]
+    : z.infer<TSchema>;
+
+type ZodInferredSafeParseResult<
+  TSchema extends ZodSchema,
+  TEnvelope extends Envelope,
+> = undefined extends TEnvelope
+  ? ParsedResult<unknown, z.infer<TSchema>>
+  : TEnvelope extends EnvelopeArrayReturnType
+    ? ParsedResult<unknown, z.infer<TSchema>>
+    : ParsedResult<unknown, z.infer<TSchema>[]>;
+
+/**
+ * The output of the parser function, can be either schema inferred type or a ParsedResult
+ */
+type ParseOutput<
+  TSchema extends ZodSchema,
+  TEnvelope extends Envelope,
+  TSafeParse = false,
+> = undefined extends TSafeParse
+  ? ZodInferredResult<TSchema, TEnvelope>
+  : TSafeParse extends true
+    ? ZodInferredSafeParseResult<TSchema, TEnvelope>
+    : TSafeParse extends false
+      ? ZodInferredResult<TSchema, TEnvelope>
+      : never;
+
 export type {
   ParserOptions,
   ParsedResult,
   ParsedResultError,
   ParsedResultSuccess,
+  ParseOutput,
 };