Zodで再帰的なスキーマを定義する関数を作ってみた

公式でも紹介されていますが、Zodで再帰的なスキーマを書くにはちょっとしたコツが必要です。
かつコードがやや読みづらくなってしまうので、なんとか一発で再帰的なスキーマを定義できないだろうか？と無駄に時間を使って考えてみました。

戦略

1. いきなり再帰的なスキーマを定義せず、自分自身を参照する部分を別の型(プレースホルダ)で置き換えたスキーマを定義する
2. 置き換えた部分を再帰化した型を作る
3. z.lazyを使って再帰的なスキーマを定義する

まずは具体例を

最初に、プレースホルダとなる型Selfを定義しておきます。他の型と区別できるように、unique symbolをキーにしたプロパティを定義してあります。

import { z } from 'zod'

const key = Symbol()
type Self = { [key]: string }

次に、z.Schema<Self>を受け取って、プレースホルダを含んだスキーマを返す関数を書きます。

function schemaBuilder(self: z.Schema<Self>) {
  return z.object({
    name: z.string(),
    child: self.optional()
  })
}

ここから型を取り出します。

type R_ = z.infer<ReturnType<typeof schemaBuilder>>
// => { name: string, child?: Self }

Rの中にあるSelfを自分自身に置き換えた型を作ります。

そこで、ReplaceSelf<T, TT>というユーティリティ型を用意しました。これは型の中にあるSelfを再帰的にTT[0]に置き換えるものです。
なぜTTがタプルなのかというと、再帰的な型のエラーを回避するためです。（なぜかはよく分かりませんが、適当に色々試していて発見しました）

ReplaceSelfは雑に作ったので、複雑な型に対応するためにもう少し手を加える必要があるかもしれません。

type ReplaceSelf<T, TT extends [unknown]> = T extends Self
  ? TT[0]
  : T extends {}
  ? { [K in keyof T]: ReplaceSelf<T[K], TT> }
  : T
  
type R = ReplaceSelf<R_, [R]>

// 確認
const r: R = { name: 'foo', child: { name: 'bar' } }

最後に、最終的なスキーマを定義しましょう。
ちょっと嘘をついてschemaBuilderにz.Schema<R>を渡せるようにして再帰的なスキーマを作ります。

const schemaBuilder_ = schemaBuilder as (self: z.Schema<R>) => z.Schema<R>
const rec = (): z.Schema<R> => schemaBuilder_(z.lazy(rec))
const schema = rec()

// 確認
schema.parse({
  name: 'foo',
  child: {
    name: 'bar',
    child: { name: 'baz' }
  }
})  // ok
schema.parse({
  name: 'foo',
  child: {}
}) // error

はい、無事に再帰的なスキーマができました！ 🎉

ユーティリティ化する

これまでやったことを関数にまとめてみます。
戻り値の型を明示的に書くことができないのが残念ですが...。

function zodRecursive<T>(builder: (self: z.Schema<Self>) => z.Schema<T>) {
  type R = ReplaceSelf<T, [R]>
  const builder_ = builder as (self: z.Schema<R>) => z.Schema<T>

  const rec = (): z.Schema<R> => builder_(z.lazy(rec))
  return rec()
}

実際に使ってみましょう。

const schema = zodRecursive((self) =>
  z.union([
    z.object({ type: z.literal('leaf'), value: z.string() }),
    z.object({ type: z.literal('branch'), children: self.array() })
  ])
)

schema.parse({
  type: 'branch',
  children: [
    { type: 'leaf', value: 'foo' },
    {
      type: 'branch',
      children: [
        { type: 'leaf', value: 'bar' }
      ]
    }
  ]
}) // ok
schema.parse({
  type: 'branch',
  children: [
    { type: 'leaf', value: 'foo' },
    { type: 'leaf', children: [] }
  ]
}) // error

はい、問題ないですね。

まとめ

• Zodで再帰的なスキーマは定義できる！
• コツは、プレースホルダを使ってスキーマを組み立てる関数を書き、そこから再帰的な型とスキーマを導くこと