Typed Zod combinator with dynamic field name

Question

My XML to JSON library emits {MyKey: T} for one-element lists, and {MyKey: T[]} for multi-element lists. The corresponding TypeScript type is type XmlJsonArray<T, element extends string> = Record<element, T | T[]>. I've used the following to implement it as a Zod schema:

const XmlJsonArray = <T, element extends string>(element: element, schema: z.Schema<T>) => {
  // TODO what is the idiomatic approach here?
  const outerSchema: Record<element, z.Schema<T | T[]>> = {} as any;
  outerSchema[element] = z.union([schema, z.array(schema)]);
  return z.object(outerSchema);
};

Is there a way to do this without using any?

Answer 1

Via @vriad:

const XmlJsonArray = <Key extends string, Schema extends z.Schema<any>>(
  element: Key,
  schema: Schema,
) => {
  return z.object({}).setKey(element, z.union([schema, z.array(schema)]));
};

const test = XmlJsonArray('asdf', z.string());

Parsing works as expected:

// both work
test.parse({ asdf: 'hello' });
test.parse({ asdf: ['hello'] });

And type inference works too: