Better egonomics for TypeScript union-narrowing typeguard function

Viewed 182

We have a utility in our codebase to generate a typeguard to narrow a discriminated union:

export type ExtractBranchFromUnion<
  UNION,
  DISCRIMINANT extends keyof UNION,
  BRANCH extends UNION[DISCRIMINANT],
> = UNION extends Record<DISCRIMINANT, BRANCH> ? UNION : never;

export function narrow<
  UNION,
  DISCRIMINANT extends keyof UNION,
  BRANCH extends UNION[DISCRIMINANT],
>(
  discriminate: DISCRIMINANT,
  branch: BRANCH,
): (
  item: UNION,
) => item is ExtractBranchFromUnion<UNION, DISCRIMINANT, BRANCH> {
  return (item): item is ExtractBranchFromUnion<UNION, DISCRIMINANT, BRANCH> =>
    item[discriminate] === branch;
}

It can be used easily in a filter call to narrow an array down to a specific union member, however it requires that you remember to add as const after a string literal branch argument, otherwise it infers the type incorrectly and gives you too wide of a result:

type A = {type: 'A'; a: string};
type B = {type: 'B'; b: string};
type SimpleUnion = A | B;
const arr: SimpleUnion[] = [];

// BAD: type is SimpleUnion[]
const badListOfBs = arr.filter(narrow('type', 'B'));

// GOOD: type is B[]
const goodListOfBs = arr.filter(narrow('type', 'B' as const));

While this issue can be partially alleviated with a premade constants,

const Types = {
  'A': 'A',
  'B': 'B',
} as const;

// OKAY: type is B[] but it requires a premade constant
const okayListOfBs = arr.filter(narrow('type', Types.B));

it still doesn't prevent the issue of people forgetting as const with a literal value and getting confused why things don't work (as well as the fact that as const just looks ugly in the example two codeblocks above).

Is there a way to get TS to automatically infer a narrower type when provided with a string literal in narrow as shown above? Or at least a way to get a nice error message to show up?

Playground link for all of the above code.

1 Answers

You just need to add extra restriction to Branch generic:

export type ExtractBranchFromUnion<
  Union,
  Discriminant extends keyof Union,
  Branch extends Union[Discriminant],
  > = Union extends Record<Discriminant, Branch> ? Union : never;

export function narrow<
  Union,
  Discriminant extends keyof Union,
  Branch extends Union[Discriminant] & PropertyKey, // <------ change is here
  >(
    discriminate: Discriminant,
    Branch: Branch,
): (
    item: Union,
  ) => item is ExtractBranchFromUnion<Union, Discriminant, Branch> {
  return (item): item is ExtractBranchFromUnion<Union, Discriminant, Branch> =>
    item[discriminate] === Branch;
}



type A = { type: 'A'; a: string };
type B = { type: 'B'; b: string };

type SimpleUnion = A | B;
const arr: SimpleUnion[] = [];

const badListOfBs = arr.filter(narrow('type', 'B')); // B[]
const goodListOfBs = arr.filter(narrow('type', 'A')); // A[]

Playground

I have added this: Branch extends Union[Discriminant] & PropertyKey.

I'd willing to bet that discriminator value should be string | number | symbol

As you might have noticed, I Capitalized all generic names because I found it hard to read. However this is the matter of preferences.

Related