Helper function to un-discriminate a discriminated union

user_11284 0
Viewed 97

TypeScript Playground Link

I'm using protobuf oneof union types generated by ts-proto, and I'd like to build a helper function that can extract a branch from the union. However, I'm unable to design a set of type constraints that expresses that the return type of my helper should be equal to the type of the corresponding property in that branch of the union. How can I make this work?

type FirstType = number;
type SecondType = string;
type ThirdType = number[];

type MyUnion  =
      { $case: 'first'; first: FirstType }
    | { $case: 'second'; second: SecondType }
    | { $case: 'third'; third: ThirdType };

type UnionCase = MyUnion['$case'];
// This doesn't work: K can't be used in the computed property
type UnionBranch<K extends UnionCase, T> = T extends {$case: K, [K]: infer Prop} ? Prop : never;

function getByCase<K extends UnionCase, T>(obj: MyUnion, $case: K): T | undefined {
  // If I were comparing against a literal, the type would be narrowed within this branch.
  if (obj.$case === $case) {
    // It can't realize that $case has the same name as the property
    return obj[$case];
  }
  return undefined;
}

const myObj = {
    $case: 'first',
    first: 5,
} as const;

const res = getByCase(myObj, 'first');
// res should be type: number
typescript
1 Answers

Both obj and obj.$case are of union types, but the compiler is not able to follow the correlation between them in order for obj[obj.$case] or the like to be accepted as-is. For example, the compiler doesn't realize that, while obj can be a { $case: 'first'; first: FirstType } and obj.$case can be "second", these cannot be both true at the same time:

function oops(obj: MyUnion) {
  return obj[obj.$case]; // error!
  // expression of type '"first" | "second" | "third"' 
  // can't be used to index type 'MyUnion'.
}

Currently the only way to get the compiler to verify type safety here is to use control flow analysis to narrow obj to each of the union members in turn, and have the compiler evaluate the expression for each narrowing. This is type safe, but redundantly repetitive and redundant:

function okay(obj: MyUnion) {
  switch (obj.$case) {
    case "first": return obj[obj.$case] // okay
    case "second": return obj[obj.$case] // okay
    case "third": return obj[obj.$case] // okay
  }
}

I filed microsoft/TypeScript#30581 about the general lack of language support for correlated union types... for now (and for the foreseeable future), you'll need to work around it.

In cases where you know something the compiler doesn't about types, you can use type assertions to tell the compiler what you know. This shifts the burden of ensuring type safety onto you, so you need to be careful that what you're telling the compiler is actually true.

Sometimes, instead of type assertions, it is easier to write an overloaded function with a single call signature. Overload implementations are more loosely checked than regular function implementations, so this has a similar effect to type assertions.

Here's how one might write your getByCase() function:

// call signature
function getByCase<T extends MyUnion, K extends UnionCase>(
  obj: T, $case: K
): K extends keyof T ? T[K] : undefined;

// implementation 
function getByCase(obj: any, $case: UnionCase) {
  if (obj.$case === $case) {
    return (obj[$case]);
  }
  return undefined;
}

The call signature represents the operation of indexing into obj (of type T) with the key $case (of type K) if it is actually a key (so type T[K]), or undefined if it is not.

The implementation is very loose, and uses any for the type of obj. You can try to make something tighter, but you will easily run afoul of the correlated union issue, and I don't know if it's worth it. Just triple check the implementation to make sure that you're doing the right thing and are not making a typo (e.g., if (obj.$case !== $case)).

Let's see if it works:

const myObj = {
  $case: 'first',
  first: 5,
} as const;

const res = getByCase(myObj, 'first'); // 5
console.log(res.toFixed(2)) // "5.00"

Looks good. res is inferred to have the type of 5, a numeric literal type, which is more specific than number. That's because myObj was declared with a const assertion. Anyway, the compiler definitely knows that res is a number, as shown by its willingness to allow res.toFixed().

Playground link to code

user_28872180
Related