Can I make Typescript narrow an object value type based on discriminated union key?

Playground link. I have a discriminated union of object types:

export type ADiscriminatedUnion = {
    type: 'a',
    dataKeyA: number,
    anotherDataKeyA: number,
} | {
    type: 'b',
    dataKeyB: number,
    anotherDataKeyB: number,
} | {
    type: 'c',
}

type AllTypes = ADiscriminatedUnion['type'];

I create an object, that maps some types (but not all) to a partial list of properties present in an object of that type:

const DiscriminatedUnionTypeToListOfDataKeysMap = {
  a: ['dataKeyA'],
  b: ['dataKeyB', 'anotherDataKeyB'],
} as const satisfies {
  [Type in AllTypes]?: Array<
    keyof Omit<Extract<ADiscriminatedUnion, { type: Type }>, 'type'>
  >;
};

type TypesWhichAreMappedToDataKeys =
  keyof typeof DiscriminatedUnionTypeToListOfDataKeysMap;

Now I want to create a function, that can only take an object of type, which was defined in the map of the second code block, and map through listed properties.

function doSomethingWithDataKeys(
  data: Extract<ADiscriminatedUnion, { type: TypesWhichAreMappedToDataKeys }>
) {
    const stringifiedDataKeys = DiscriminatedUnionTypeToListOfDataKeysMap[data.type]
        .map(key => `${data[key]}`).join('_');
}

This errors on data[key]:

Property 'dataKeyA' does not exist on type '{ type: "a"; dataKeyA: number; anotherDataKeyA: number; } | { type: "b"; dataKeyB: number; anotherDataKeyB: number; }'.

It seems like TypeScript doesn't narrow DiscriminatedUnionTypeToListOfDataKeysMap[data.type] to match the resulting value (array of property names) with a passed object's properties. How do I instruct it to do so?

Solution

You've run into a variant of microsoft/TypeScript#30581. TypeScript doesn't understand the correlation between data.type and data[key] when data is of a union type. It would need to do so "at once", it doesn't preemptively narrow data to each union member of ADiscriminatedUnion and check each narrowing separately. That would be okay for a two-member union happening in one place, but that sort of thing does not scale.

The way to make TypeScript analyze something general at once is to use generics. For correlated unions there is a particular sort of refactoring described in microsoft/TypeScript#47109 which works. The approach is to represent everything you're doing in terms of some "base" key-value type, or in terms of mapped types over that type, or in terms of generic indexes into those types.

For your example the cleanest refactoring is as follows. First, let's define the base key-value type:

interface BaseType {
  a: {
    dataKeyA: number;
    anotherDataKeyA: number;
  };
  b: {
    dataKeyB: number;
    anotherDataKeyB: number;
  };
  c: {
    dataKeyC: number;
    anotherDataKeyC: number;
  };
}

Then your discriminated union can be made generic in BaseType, and in particular, will be a distributive object type as coined in ms/TS#47109:

type ADiscriminatedUnion<K extends keyof BaseType = keyof BaseType> =
  { [P in K]: { type: P } & BaseType[P] }[K]

The type named ADiscriminatedUnion with no type argument (which uses the default type argument of keyof BaseType) is equivalent to your version, but it is represented in a way as to make the rest of the code easier for the compiler to follow:

export const DiscriminatedUnionTypeToListOfDataKeysMap = {
  a: ['dataKeyA'],
  b: ['dataKeyB', 'anotherDataKeyB'],
} as const satisfies { [K in keyof BaseType]?: (keyof BaseType[K])[] }

type TypesWhichAreMappedToDataKeys = 
  keyof typeof DiscriminatedUnionTypeToListOfDataKeysMap;

function doSomethingWithDataKeys<K extends TypesWhichAreMappedToDataKeys>(
  data: ADiscriminatedUnion<K>
) {
  const d: { [K in TypesWhichAreMappedToDataKeys]: (keyof BaseType[K])[] } =
    DiscriminatedUnionTypeToListOfDataKeysMap;
  const stringifiedDataKeys = d[data.type]
    .map(key => `${data[key]}`).join('_');
}

I had to widen DiscriminatedUnionTypeToListOfDataKeysMap to the type { [K in TypesWhichAreMappedToDataKeys]: (keyof BaseType[K])[] }, which lets the compiler understand that DiscriminatedUnionTypeToListOfDataKeysMap[data.type] is of type (keyof BaseType[K])[], which means that key is of type keyof BaseType[K], which works, because data is assignable to BaseType[K], and so data[key] is BaseType[K][keyof BaseType[K]], which is the union of all BaseType[K]'s value types. And since all of those are the serializable number, the code compiles.

Note that you could start with your original ADiscriminatedUnion type on use it to compute BaseType, but that's more complicated:

type BaseType = { [T in ADiscriminatedUnion as T['type']]: 
  { [K in keyof T as K extends "type" ? never : K]: T[K] } 
};

And then you'd need to use the distributive object type version of ADiscriminatedUnion instead of your original one, for TypeScript to follow the correlation:

type DiscU<K extends keyof BaseType = keyof BaseType> = 
  { [P in K]: { type: P } & BaseType[P] }[K]

I'd only go this route if you have no control over the original ADiscriminatedUnion type. But either way it works.

Playground link to code