discriminated union matching

index.ts

type AnyVariant = { type: string };

type PayloadOf<U extends AnyVariant, K extends U["type"]> =
  Extract<U, { type: K }> extends infer Member
    ? Member[K & keyof Member]
    : never;

type Handler<U extends AnyVariant, R> = {
  [K in U["type"]]: (payload: PayloadOf<U, K>) => R;
};

export function match<U extends AnyVariant, H extends Handler<U, unknown>>(
  variant: U,
  handlers: H,
  fallback?: (variant: U) => ReturnType<H[U["type"]]>,
): ReturnType<H[U["type"]]> {
  const { type } = variant;

  if (type in handlers && type in variant) {
    const payload = (variant as any)[type];
    return (handlers as any)[type](payload);
  }

  if (fallback) {
    return fallback(variant);
  }

  throw new Error(`unreachable`);
}

type PartialHandlerReturn<
  U extends AnyVariant,
  H extends Partial<Handler<U, unknown>>,
> = {
  [K in U["type"]]: K extends keyof H
    ? H[K] extends (...args: any) => infer R
      ? R
      : never
    : never;
}[U["type"]];

export function matchPartial<
  U extends AnyVariant,
  H extends Partial<Handler<U, unknown>>,
  F = never,
>(
  variant: U,
  handlers: H,
  fallback?: (variant: U) => F,
): PartialHandlerReturn<U, H> | F {
  const defaultFallback = ({type}: U): never => {
    throw new Error(`matchPartial: no handler for variant type "${type}"`);
  };
  return match(
    variant,
    handlers as Handler<U, unknown>,
    (fallback ?? defaultFallback) as any,
  ) as PartialHandlerReturn<U, H> | F;
}

// Usage:

type MyVariant =
  | {
      type: "foo";
      foo: string;
    }
  | {
      type: "bar";
      bar: number;
    };

export const handleMyVariant = (variant: MyVariant): number => {
  const num = match(variant, {
    foo: (v) => parseFloat(v),
    bar: (v) => v,
  });
  return num;
};

export const handleMyVariantPartial = (variant: MyVariant): number => {
  const num = matchPartial(variant, {
    foo: (v) => parseFloat(v),
  });
  return num;
};