Complete zod plugin with proprietary brands

example/endpoints/accept-raw.ts

@@ -5,9 +5,10 @@ import { taggedEndpointsFactory } from "../factories";
 export const rawAcceptingEndpoint = taggedEndpointsFactory.build({
   method: "post",
   tag: "files",
-  input: ez
-    .raw() // requires to enable rawParser option in server config
-    .extend({}), // additional inputs, route params for example, if needed
+  // requires to enable rawParser option in server config:
+  input: ez.raw({
+    /* additional inputs, route params for example, if needed */
+  }),
   output: z.object({ length: z.number().int().nonnegative() }),
   handler: async ({ input: { raw } }) => ({
     length: raw.length, // input.raw is populated automatically when rawParser is set in config

src/date-in-schema.ts

@@ -1,8 +1,7 @@
 import { z } from "zod";
-import { proprietary } from "./metadata";
 import { isValidDate } from "./schema-helpers";
 
-export const ezDateInKind = "DateIn";
+export const ezDateInBrand = Symbol("DateIn");
 
 export const dateIn = () => {
   const schema = z.union([
@@ -11,8 +10,10 @@ export const dateIn = () => {
     z.string().datetime({ local: true }),
   ]);
 
-  return proprietary(
-    ezDateInKind,
-    schema.transform((str) => new Date(str)).pipe(z.date().refine(isValidDate)),
-  );
+  return schema
+    .transform((str) => new Date(str))
+    .pipe(z.date().refine(isValidDate))
+    .brand(ezDateInBrand);
 };
+
+export type DateInSchema = ReturnType<typeof dateIn>;

src/date-out-schema.ts

@@ -1,14 +1,13 @@
 import { z } from "zod";
-import { proprietary } from "./metadata";
 import { isValidDate } from "./schema-helpers";
 
-export const ezDateOutKind = "DateOut";
+export const ezDateOutBrand = Symbol("DateOut");
 
 export const dateOut = () =>
-  proprietary(
-    ezDateOutKind,
-    z
-      .date()
-      .refine(isValidDate)
-      .transform((date) => date.toISOString()),
-  );
+  z
+    .date()
+    .refine(isValidDate)
+    .transform((date) => date.toISOString())
+    .brand(ezDateOutBrand);
+
+export type DateOutSchema = ReturnType<typeof dateOut>;

src/deep-checks.ts

@@ -1,9 +1,9 @@
 import { z } from "zod";
 import { IOSchema } from "./io-schema";
 import { isProprietary } from "./metadata";
-import { ezRawKind } from "./raw-schema";
+import { ezRawBrand } from "./raw-schema";
 import { HandlingRules, SchemaHandler } from "./schema-walker";
-import { ezUploadKind } from "./upload-schema";
+import { ezUploadBrand } from "./upload-schema";
 
 /** @desc Check is a schema handling rule returning boolean */
 type Check<T extends z.ZodTypeAny> = SchemaHandler<T, boolean>;
@@ -95,12 +95,12 @@ export const hasTransformationOnTop = (subject: IOSchema): boolean =>
 export const hasUpload = (subject: IOSchema) =>
   hasNestedSchema({
     subject,
-    condition: (schema) => isProprietary(schema, ezUploadKind),
+    condition: (schema) => isProprietary(schema, ezUploadBrand),
   });
 
 export const hasRaw = (subject: IOSchema) =>
   hasNestedSchema({
     subject,
-    condition: (schema) => isProprietary(schema, ezRawKind),
+    condition: (schema) => isProprietary(schema, ezRawBrand),
     maxDepth: 3,
   });

src/documentation-helpers.ts

@@ -52,10 +52,10 @@ import {
   ucFirst,
 } from "./common-helpers";
 import { InputSource, TagsConfig } from "./config-type";
-import { ezDateInKind } from "./date-in-schema";
-import { ezDateOutKind } from "./date-out-schema";
+import { DateInSchema, ezDateInBrand } from "./date-in-schema";
+import { DateOutSchema, ezDateOutBrand } from "./date-out-schema";
 import { DocumentationError } from "./errors";
-import { ezFileKind } from "./file-schema";
+import { FileSchema, ezFileBrand } from "./file-schema";
 import { IOSchema } from "./io-schema";
 import {
   LogicalContainer,
@@ -64,15 +64,15 @@ import {
 } from "./logical-container";
 import { getMeta } from "./metadata";
 import { Method } from "./method";
-import { RawSchema, ezRawKind } from "./raw-schema";
+import { RawSchema, ezRawBrand } from "./raw-schema";
 import {
   HandlingRules,
   HandlingVariant,
   SchemaHandler,
   walkSchema,
 } from "./schema-walker";
 import { Security } from "./security";
-import { ezUploadKind } from "./upload-schema";
+import { UploadSchema, ezUploadBrand } from "./upload-schema";
 
 /* eslint-disable @typescript-eslint/no-use-before-define */
 
@@ -146,7 +146,7 @@ export const depictAny: Depicter<z.ZodAny> = () => ({
   format: "any",
 });
 
-export const depictUpload: Depicter<z.ZodType> = (ctx) => {
+export const depictUpload: Depicter<UploadSchema> = (ctx) => {
   assert(
     !ctx.isResponse,
     new DocumentationError({
@@ -160,15 +160,18 @@ export const depictUpload: Depicter<z.ZodType> = (ctx) => {
   };
 };
 
-export const depictFile: Depicter<z.ZodType> = ({ schema }) => ({
-  type: "string",
-  format:
-    schema instanceof z.ZodString
-      ? schema._def.checks.find((check) => check.kind === "base64")
-        ? "byte"
-        : "file"
-      : "binary",
-});
+export const depictFile: Depicter<FileSchema> = ({ schema }) => {
+  const subject = schema.unwrap();
+  return {
+    type: "string",
+    format:
+      subject instanceof z.ZodString
+        ? subject._def.checks.find((check) => check.kind === "base64")
+          ? "byte"
+          : "file"
+        : "binary",
+  };
+};
 
 export const depictUnion: Depicter<z.ZodUnion<z.ZodUnionOptions>> = ({
   schema: { options },
@@ -317,7 +320,7 @@ export const depictObject: Depicter<z.ZodObject<z.ZodRawShape>> = ({
  * */
 export const depictNull: Depicter<z.ZodNull> = () => ({ type: "null" });
 
-export const depictDateIn: Depicter<z.ZodType> = (ctx) => {
+export const depictDateIn: Depicter<DateInSchema> = (ctx) => {
   assert(
     !ctx.isResponse,
     new DocumentationError({
@@ -336,7 +339,7 @@ export const depictDateIn: Depicter<z.ZodType> = (ctx) => {
   };
 };
 
-export const depictDateOut: Depicter<z.ZodType> = (ctx) => {
+export const depictDateOut: Depicter<DateOutSchema> = (ctx) => {
   assert(
     ctx.isResponse,
     new DocumentationError({
@@ -628,7 +631,7 @@ export const depictLazy: Depicter<z.ZodLazy<z.ZodTypeAny>> = ({
 };
 
 export const depictRaw: Depicter<RawSchema> = ({ next, schema }) =>
-  next(schema.shape.raw);
+  next(schema.unwrap().shape.raw);
 
 const enumerateExamples = (examples: unknown[]): ExamplesObject | undefined =>
   examples.length
@@ -669,6 +672,9 @@ export const extractObjectSchema = (
   if (subject instanceof z.ZodObject) {
     return subject;
   }
+  if (subject instanceof z.ZodBranded) {
+    return extractObjectSchema(subject.unwrap(), tfError);
+  }
   if (
     subject instanceof z.ZodUnion ||
     subject instanceof z.ZodDiscriminatedUnion
@@ -779,11 +785,11 @@ export const depicters: HandlingRules<
   ZodPipeline: depictPipeline,
   ZodLazy: depictLazy,
   ZodReadonly: depictReadonly,
-  [ezFileKind]: depictFile,
-  [ezUploadKind]: depictUpload,
-  [ezDateOutKind]: depictDateOut,
-  [ezDateInKind]: depictDateIn,
-  [ezRawKind]: depictRaw,
+  [ezFileBrand]: depictFile,
+  [ezUploadBrand]: depictUpload,
+  [ezDateOutBrand]: depictDateOut,
+  [ezDateInBrand]: depictDateIn,
+  [ezRawBrand]: depictRaw,
 };
 
 export const onEach: Depicter<z.ZodTypeAny, "each"> = ({

src/file-schema.ts

@@ -1,17 +1,16 @@
 import { z } from "zod";
-import { proprietary } from "./metadata";
 
-export const ezFileKind = "File";
+export const ezFileBrand = Symbol("File");
 
 const bufferSchema = z.custom<Buffer>((subject) => Buffer.isBuffer(subject), {
   message: "Expected Buffer",
 });
 
 const variants = {
-  buffer: () => proprietary(ezFileKind, bufferSchema),
-  string: () => proprietary(ezFileKind, z.string()),
-  binary: () => proprietary(ezFileKind, bufferSchema.or(z.string())),
-  base64: () => proprietary(ezFileKind, z.string().base64()),
+  buffer: () => bufferSchema.brand(ezFileBrand),
+  string: () => z.string().brand(ezFileBrand),
+  binary: () => bufferSchema.or(z.string()).brand(ezFileBrand),
+  base64: () => z.string().base64().brand(ezFileBrand),
 };
 
 type Variants = typeof variants;
@@ -22,3 +21,5 @@ export function file<K extends Variant>(variant: K): ReturnType<Variants[K]>;
 export function file<K extends Variant>(variant?: K) {
   return variants[variant || "string"]();
 }
+
+export type FileSchema = ReturnType<typeof file>;

src/io-schema.ts

@@ -1,6 +1,7 @@
 import { z } from "zod";
 import { copyMeta } from "./metadata";
 import { AnyMiddlewareDef } from "./middleware";
+import { RawSchema } from "./raw-schema";
 
 type Refined<T extends z.ZodTypeAny> =
   T extends z.ZodType<infer O> ? z.ZodEffects<T | Refined<T>, O, O> : never;
@@ -14,7 +15,8 @@ export type IOSchema<U extends z.UnknownKeysParam = any> =
   | z.ZodUnion<[IOSchema<U>, ...IOSchema<U>[]]>
   | z.ZodIntersection<IOSchema<U>, IOSchema<U>>
   | z.ZodDiscriminatedUnion<string, z.ZodObject<any, U>[]>
-  | Refined<z.ZodObject<any, U>>;
+  | Refined<z.ZodObject<any, U>>
+  | RawSchema;
 
 export type ProbableIntersection<
   A extends IOSchema<"strip"> | null,

src/metadata.ts

@@ -1,15 +1,15 @@
 import { combinations, isObject } from "./common-helpers";
 import { z } from "zod";
 import { clone, mergeDeepRight } from "ramda";
-import { ProprietaryKind } from "./proprietary-schemas";
+import { ProprietaryBrand } from "./proprietary-schemas";
 
 export const metaSymbol = Symbol.for("express-zod-api");
 
 export interface Metadata<T extends z.ZodTypeAny> {
-  kind?: ProprietaryKind;
   examples: z.input<T>[];
   /** @override ZodDefault::_def.defaultValue() in depictDefault */
   defaultLabel?: string;
+  brand?: string | number | symbol;
 }
 
 declare module "zod" {
@@ -53,6 +53,16 @@ const defaultLabeler = function (
   return copy;
 };
 
+const brander = function (this: z.ZodType, brand?: string | number | symbol) {
+  return new z.ZodBranded({
+    typeName: z.ZodFirstPartyTypeKind.ZodBranded,
+    type: this,
+    description: this._def.description,
+    errorMap: this._def.errorMap,
+    [metaSymbol]: { examples: [], ...clone(this._def[metaSymbol]), brand },
+  });
+};
+
 /** @see https://github.com/colinhacks/zod/blob/90efe7fa6135119224412c7081bd12ef0bccef26/plugin/effect/src/index.ts#L21-L31 */
 if (!(metaSymbol in globalThis)) {
   (globalThis as Record<symbol, unknown>)[metaSymbol] = true;
@@ -65,6 +75,16 @@ if (!(metaSymbol in globalThis)) {
       },
     },
   );
+  Object.defineProperty(
+    z.ZodType.prototype,
+    "brand" satisfies keyof z.ZodType,
+    {
+      set() {}, // this is required to override the existing method
+      get(): z.ZodType["brand"] {
+        return brander.bind(this);
+      },
+    },
+  );
   Object.defineProperty(
     z.ZodDefault.prototype,
     "label" satisfies keyof z.ZodDefault<z.ZodTypeAny>,
@@ -104,14 +124,5 @@ export const copyMeta = <A extends z.ZodTypeAny, B extends z.ZodTypeAny>(
   return result;
 };
 
-export const proprietary = <T extends z.ZodTypeAny>(
-  kind: ProprietaryKind,
-  subject: T,
-) => {
-  const schema = cloneSchema(subject);
-  schema._def[metaSymbol].kind = kind;
-  return schema;
-};
-
-export const isProprietary = (schema: z.ZodTypeAny, kind: ProprietaryKind) =>
-  getMeta(schema, "kind") === kind;
+export const isProprietary = (schema: z.ZodTypeAny, brand: ProprietaryBrand) =>
+  getMeta(schema, "brand") === brand;

src/proprietary-schemas.ts

@@ -1,14 +1,14 @@
-import { dateIn, ezDateInKind } from "./date-in-schema";
-import { dateOut, ezDateOutKind } from "./date-out-schema";
-import { ezFileKind, file } from "./file-schema";
-import { ezRawKind, raw } from "./raw-schema";
-import { ezUploadKind, upload } from "./upload-schema";
+import { dateIn, ezDateInBrand } from "./date-in-schema";
+import { dateOut, ezDateOutBrand } from "./date-out-schema";
+import { ezFileBrand, file } from "./file-schema";
+import { ezRawBrand, raw } from "./raw-schema";
+import { ezUploadBrand, upload } from "./upload-schema";
 
 export const ez = { dateIn, dateOut, file, upload, raw };
 
-export type ProprietaryKind =
-  | typeof ezFileKind
-  | typeof ezDateInKind
-  | typeof ezDateOutKind
-  | typeof ezUploadKind
-  | typeof ezRawKind;
+export type ProprietaryBrand =
+  | typeof ezFileBrand
+  | typeof ezDateInBrand
+  | typeof ezDateOutBrand
+  | typeof ezUploadBrand
+  | typeof ezRawBrand;

src/raw-schema.ts

@@ -1,11 +1,13 @@
 import { z } from "zod";
 import { file } from "./file-schema";
-import { proprietary } from "./metadata";
 
-export const ezRawKind = "Raw";
+export const ezRawBrand = Symbol("Raw");
 
 /** Shorthand for z.object({ raw: ez.file("buffer") }) */
-export const raw = () =>
-  proprietary(ezRawKind, z.object({ raw: file("buffer") }));
+export const raw = <S extends z.ZodRawShape>(extra: S = {} as S) =>
+  z
+    .object({ raw: file("buffer") })
+    .extend(extra)
+    .brand(ezRawBrand);
 
 export type RawSchema = ReturnType<typeof raw>;