Conditional keys based on value of another key with Zod

Question

I'm making a project with the TMDB API and trying to make it super type-safe to reinforce some of the TypeScript stuff I'm learning. I'm using Zod to describe the shape of the data returned by the API.

However, I've noticed that depending on the request parameters, the API can send back data with different keys. Specifically, if the API is sending back data from the "trending" endpoint where data.media_type = "movie" it also has the keys title, original_title, and release_date. But if data.media_type = "tv", those three keys are renamed name, original_name, and first_air_date, respectively, as well as a new key of origin_country being added.

As a result, I described the shape of my data like this:

const mediaType = ["all", "movie", "tv", "person"] as const

const dataShape = z.object({
    page: z.number(),
    results: z.array(z.object({
        adult: z.boolean(),
        backdrop_path: z.string(),
        first_air_date: z.string().optional(),
        release_date: z.string().optional(),
        genre_ids: z.array(z.number()),
        id: z.number(),
        media_type: z.enum(mediaType),
        name: z.string().optional(),
        title: z.string().optional(),
        origin_country: z.array(z.string()).optional(),
        original_language: z.string().default("en"),
        original_name: z.string().optional(),
        original_title: z.string().optional(),
        overview: z.string(),
        popularity: z.number(),
        poster_path: z.string(),
        vote_average: z.number(),
        vote_count: z.number()
    })),
    total_pages: z.number(),
    total_results: z.number()
})

Basically, I've added .optional() to every troublesome key. Obviously, this isn't very type-safe. Is there a way to specify that the origin_country key only exists when media_type is equal to tv, or that the key name or title are both a z.string(), but whose existence is conditional?

It may be worth stating that the media_type is also specified outside of the returned data, specifically in the input to the API call (which for completeness looks like this, using tRPC):

import { tmdbRoute } from "../utils"
import { publicProcedure } from "../trpc"

export const getTrending = publicProcedure
    .input(z.object({
        mediaType: z.enum(mediaType).default("all"),
        timeWindow: z.enum(["day", "week"]).default("day")
    }))
    .output(dataShape)
    .query(async ({ input }) => {
        return await fetch(tmdbRoute(`/trending/${input.mediaType}/${input.timeWindow}`))
            .then(res => res.json())
    })

Any help is appreciated!

Edit: I have learned about the Zod method of discriminatedUnion() since posting this, but if that's the correct approach I'm struggling to implement it. Currently have something like this:

const indiscriminateDataShape = z.object({
    page: z.number(),
    results: z.array(
        z.object({
            adult: z.boolean(),
            backdrop_path: z.string(),
            genre_ids: z.array(z.number()),
            id: z.number(),
            media_type: z.enum(mediaType),
            original_language: z.string().default("en"),
            overview: z.string(),
            popularity: z.number(),
            poster_path: z.string(),
            vote_average: z.number(),
            vote_count: z.number()
        })
    ),
    total_pages: z.number(),
    total_results: z.number()
})

const dataShape = z.discriminatedUnion('media_type', [
    z.object({
        media_type: z.literal("tv"),
        name: z.string(),
        first_air_date: z.string(),
        original_name: z.string(),
        origin_country: z.array(z.string())
    }).merge(indiscriminateDataShape),
    z.object({
        media_type: z.literal("movie"),
        title: z.string(),
        release_date: z.string(),
        original_title: z.string()
    }).merge(indiscriminateDataShape),
    z.object({
        media_type: z.literal("all")
    }).merge(indiscriminateDataShape),
    z.object({
        media_type: z.literal("person")
    }).merge(indiscriminateDataShape)
])

Making the request with any value for media_type with the above code logs the error "Invalid discriminator value. Expected 'tv' | 'movie' | 'all' | 'person'"

Answer 1

It's a great example of using Zod to validate schemas. Discriminated unions are the solution to your problem as you noticed, but I think that is a misunderstood of the API schema in your last implementation.

Making some requests to TMDB API, the most basic schema is something like this:

const schema = {
  page: 1,
  results: [],
  total_pages: 100,
  total_results: 200,
}

So, in your Zod schema, you need to consider that first. After, we will use the z.discriminatedUnion() function inside results property. I'm also considering the merge or extending baseShape in the last step (after discriminatedUnion).

const baseShape = z.object({
  adult: z.boolean(),
  backdrop_path: z.string(),
  genre_ids: z.array(z.number()),
  id: z.number(),
  original_language: z.string().default('en'),
  overview: z.string(),
  popularity: z.number(),
  poster_path: z.string(),
  vote_average: z.number(),
  vote_count: z.number(),
});

const resultShape = z
  .discriminatedUnion('media_type', [
    // tv shape
    z.object({
      media_type: z.literal('tv'),
      name: z.string(),
      first_air_date: z.string(),
      original_name: z.string(),
      origin_country: z.array(z.string()),
    }),

    // movie shape
    z.object({
      media_type: z.literal('movie'),
      title: z.string(),
      release_date: z.string(),
      original_title: z.string(),
    }),

    // all shape
    z.object({
      media_type: z.literal('all'),
    }),
  ])
  .and(baseShape);

const requestShape = z.object({
  page: z.number(),
  results: z.array(resultShape),
  total_pages: z.number(),
  total_results: z.number(),
});

You can see the full implementation here in StackBlitz with some data to test.