Narrow discriminated union by function parameter-CodePudding

I have a store which keys a few different types by their IDs

interface Dog { type: "dog"; woofs: string; }
interface Cat { type: "cat"; meows: string; }
type Pet = Dog | Cat;
type AnimalState = Record<string, Pet>

And I would like a function that returns a pet specific type, by it's key/ID.

function getSpecificAnimal(state: AnimalState, key: string, type: Pet["type"]) {
  const pet = state[key];

  if (pet.type === type) {
    return pet;
  }

  throw new Error("pet was wrong type");
}

const aDog = getSpecificAnimal(state, "a", "dog"); // should return type Dog, or throw an exception

But I cannot figure out how to get Typescript to narrow down the return of the getSpecificAnimal function.

I tried using generics and a type map to relate all the arguments, but Typescript still was not satisfied

type PetTypeMap = {
  dog: Dog,
  cat: Cat,
}
function getAnimalOfType<T extends Pet, K extends T["type"]>(animals: AnimalState, id: string, type: K): PetTypeMap[K] {
  const pet = animals[id];
  //    ^?

  if (pet.type === type) {
    return pet;
  }

  throw new Error("wrong pet type");
}

I've tried a bunch of other ways and just cannot get this to work. Is this type of narrowing not possible with Typescript?

CodePudding user response：

It's not possible to implement getAnimalOfType() with compiler-verified type safety as of TypeScript 4.7. In order for the return type of depend on its inputs, you'd either need to make it an overloaded function or a generic one. Callers will be happy either way, but the compiler can't help you much inside the implementation.

Overloads are easy enough to write and call:

function getSpecificAnimal(state: AnimalState, key: string, type: "dog"): Dog;
function getSpecificAnimal(state: AnimalState, key: string, type: "cat"): Cat;
function getSpecificAnimal(state: AnimalState, key: string, type: Pet["type"]) {
    const pet = state[key];

    if (pet.type === type) {
        return pet; // no error
    }

    throw new Error("pet was wrong type");
}

const state = {
    a: { type: "dog", woofs: "a lot" },
    b: { type: "cat", meows: "incessantly" }
} as const;
const aDog = getSpecificAnimal(state, "a", "dog");
// const aDog: Dog
console.log(aDog.woofs.toUpperCase()) // A LOT

but their implementations are loosely checked (see microsoft/TypeScript#10765). So even though the compiler won't complain, you need to take care that you've implemented it correctly:

if (pet.type !== type) { // <-- wrong check
    return pet; // <-- still no error
}

The generic version would have a call signature like

declare function getSpecificAnimal<K extends Pet['type']>(
    state: AnimalState, key: string, type: K
): Extract<Pet, { type: K }>;

using the Extract<T, U> utility type in its return type to say that you want to get just the member of the Pet union whose type property is the same as that of the type parameter. It is also easy for callers:

const state = {
    a: { type: "dog", woofs: "a lot" },
    b: { type: "cat", meows: "incessantly" }
} as const;
const aDog = getSpecificAnimal(state, "a", "dog");
// const aDog: Dog
console.log(aDog.woofs.toUpperCase()) // A LOT

But implementation will give you the opposite problem from overloads; you'll get compiler warnings even if you write it correctly. Generics and narrowing don't often play nicely together:

function getSpecificAnimal<K extends Pet['type']>(
    state: AnimalState, key: string, type: K
): Extract<Pet, { type: K }> {
    const pet = state[key];

    if (pet.type === type) {
        return pet; // error,
        // Type 'Pet' is not assignable to type 
        // 'Extract<Dog, { type: K; }> | Extract<Cat, { type: K; }>'
    }

    throw new Error("pet was wrong type");
}

There are a lot of open issues about it in GitHub. For the specific flavor of generic narrowing you'd need, there's microsoft/TypeScript#46899. It's listed as "Awaiting More Feedback", meaning they want to hear compelling use cases from the community before considering implementing it. You might want to go there and contribute, but I don't know that anything would ever come of it.

For now the easiest way to proceed is just to use a type assertion, which is the general-purpose recourse for situations where you know more about the type of a value than the compiler does:

function getSpecificAnimal<K extends Pet['type']>(state: AnimalState, key: string, type: K) {
    const pet = state[key];

    if (pet.type === type) {
        return pet as Extract<Pet, { type: K }>; // <-- assert here
    }

    throw new Error("pet was wrong type");
}

But this is back to the same issue as with the overload implementation; you have to take care to implement correctly, because you won't get an error if you do it wrong:

if (pet.type !== type) { // <-- wrong check
    return pet as Extract<Pet, { type: K }>; // <-- still no error
}

So those are the options. You can overload, or you can use a generic with a conditional return type. Callers are mostly happy, but as implementer you need to take care because the compiler can't do much to help.

Playground link to code