feat: allow to append or override schema via jsdoc tag (#266)

implementation taken from #125

README.md

@@ -150,11 +150,13 @@ export const heroContactSchema = z.object({
 
 Other JSDoc tags are available:
 
-| JSDoc keyword          | JSDoc Example            | Description                               | Generated Zod                      |
-| ---------------------- | ------------------------ | ----------------------------------------- | ---------------------------------- |
-| `@description {value}` | `@description Full name` | Sets the description of the property      | `z.string().describe("Full name")` |
-| `@default {value}`     | `@default 42`            | Sets a default value for the property     | `z.number().default(42)`           |
-| `@strict`              | `@strict`                | Adds the `strict()` modifier to an object | `z.object().strict()`              |
+| JSDoc keyword          | JSDoc Example            | Description                                                                                                                               | Generated Zod                      |
+|------------------------|--------------------------|-------------------------------------------------------------------------------------------------------------------------------------------|------------------------------------|
+| `@description {value}` | `@description Full name` | Sets the description of the property                                                                                                      | `z.string().describe("Full name")` |
+| `@default {value}`     | `@default 42`            | Sets a default value for the property                                                                                                     | `z.number().default(42)`           |
+| `@strict`              | `@strict`                | Adds the `strict()` modifier to an object                                                                                                 | `z.object().strict()`              |
+| `@schema`              | `@schema .catch('foo')`  | If value starts with a `.`, appends the specified value to the generated schema. Otherwise this value will override the generated schema. | `z.string().catch('foo')`          |
+  |
 
 ## JSDoc tags for elements of `string` and `number` arrays
 

src/core/generateZodSchema.test.ts

@@ -964,6 +964,48 @@ describe("generateZodSchema", () => {
     `);
   });
 
+  it("should append schema based on `schema` tag", () => {
+    const source = `export interface HeroContact {
+      /**
+       * The email of the hero.
+       *
+       * @schema .trim().catch('[email protected]')
+       */
+      email: string;
+    }`;
+    expect(generate(source)).toMatchInlineSnapshot(`
+      "export const heroContactSchema = z.object({
+          /**
+           * The email of the hero.
+           *
+           * @schema .trim().catch('[email protected]')
+           */
+          email: z.string().trim().catch('[email protected]')
+      });"
+    `);
+  });
+
+  it("should overrride schema based on `schema` tag", () => {
+    const source = `export interface HeroContact {
+      /**
+       * The email of the hero.
+       *
+       * @schema coerce.int()
+       */
+      age: number;
+    }`;
+    expect(generate(source)).toMatchInlineSnapshot(`
+      "export const heroContactSchema = z.object({
+          /**
+           * The email of the hero.
+           *
+           * @schema coerce.int()
+           */
+          age: z.coerce.int()
+      });"
+    `);
+  });
+
   it("should generate custom error message for `format` tag", () => {
     const source = `export interface HeroContact {
       /**

src/core/generateZodSchema.ts

@@ -62,6 +62,21 @@ type SchemaExtensionClause = {
   omitOrPickKeys?: ts.TypeNode;
 };
 
+interface BuildZodPrimitiveParams {
+  z: string;
+  typeNode: ts.TypeNode;
+  isOptional: boolean;
+  isNullable?: boolean;
+  isPartial?: boolean;
+  isRequired?: boolean;
+  jsDocTags: JSDocTags;
+  customJSDocFormatTypes: CustomJSDocFormatTypes;
+  sourceFile: ts.SourceFile;
+  dependencies: string[];
+  getDependencyName: (identifierName: string) => string;
+  skipParseJSDoc: boolean;
+}
+
 /**
  * Generate zod schema declaration
  *
@@ -288,6 +303,35 @@ function buildZodProperties({
 }
 
 function buildZodPrimitive({
+  jsDocTags,
+  z,
+  ...rest
+}: BuildZodPrimitiveParams):
+  | ts.CallExpression
+  | ts.Identifier
+  | ts.PropertyAccessExpression {
+  const schema = jsDocTags.schema;
+  delete jsDocTags.schema;
+  const generatedSchema = buildZodPrimitiveInternal({ jsDocTags, z, ...rest });
+  // schema not specified? return generated one
+  if (!schema) {
+    return generatedSchema;
+  }
+  // schema starts with dot? append it
+  if (schema.startsWith(".")) {
+    return f.createPropertyAccessExpression(
+      generatedSchema,
+      f.createIdentifier(schema.slice(1))
+    );
+  }
+  // otherwise use provided schema verbatim
+  return f.createPropertyAccessExpression(
+    f.createIdentifier(z),
+    f.createIdentifier(schema)
+  );
+}
+
+function buildZodPrimitiveInternal({
   z,
   typeNode,
   isOptional,
@@ -300,20 +344,10 @@ function buildZodPrimitive({
   dependencies,
   getDependencyName,
   skipParseJSDoc,
-}: {
-  z: string;
-  typeNode: ts.TypeNode;
-  isOptional: boolean;
-  isNullable?: boolean;
-  isPartial?: boolean;
-  isRequired?: boolean;
-  jsDocTags: JSDocTags;
-  customJSDocFormatTypes: CustomJSDocFormatTypes;
-  sourceFile: ts.SourceFile;
-  dependencies: string[];
-  getDependencyName: (identifierName: string) => string;
-  skipParseJSDoc: boolean;
-}): ts.CallExpression | ts.Identifier | ts.PropertyAccessExpression {
+}: BuildZodPrimitiveParams):
+  | ts.CallExpression
+  | ts.Identifier
+  | ts.PropertyAccessExpression {
   const zodProperties = jsDocTagToZodProperties(
     jsDocTags,
     customJSDocFormatTypes,

src/core/jsDocTags.ts

@@ -68,6 +68,7 @@ export interface JSDocTagsBase {
    */
   pattern?: string;
   strict?: boolean;
+  schema?: string;
 }
 
 export type ElementJSDocTags = Pick<
@@ -94,6 +95,7 @@ const jsDocTagKeys: Array<keyof JSDocTags> = [
   "maxLength",
   "format",
   "pattern",
+  "schema",
   "elementDescription",
   "elementMinimum",
   "elementMaximum",
@@ -177,6 +179,7 @@ export function getJSDocTags(nodeType: ts.Node, sourceFile: ts.SourceFile) {
             break;
           case "description":
           case "elementDescription":
+          case "schema":
           case "pattern":
           case "elementPattern":
             if (tag.comment) {