How to infer the exact return type from a heterogeneous getter map?

I have a strongly typed State object and a separate set of ParamKeys. Each StateKey corresponds to one or more ParamKeys.

I want to build a params map where each ParamKey value is derived from State, while preserving the exact value types — no any types.

Example

const ParamKey = { A: 'a', B1: 'b1', C1: 'c1', C2: 'c2' } as const;
type ParamKey = (typeof ParamKey)[keyof typeof ParamKey];
type ParamValue = Primitive | Primitive[];

const StateKey = { A: 'a', B: 'b', C: 'c' } as const;
type StateKey = (typeof StateKey)[keyof typeof StateKey];

type State = {
    [StateKey.A]: number;
    [StateKey.B]: string[];
    [StateKey.C]: { key1: string; key2: number };
};
const state: State = {
    [StateKey.A]: 100,
    [StateKey.B]: ['srt0', 'srt1'],
    [StateKey.C]: { key1: 'value', key2: 2000 },
};

Each StateKey has a corresponding getter that maps it to one or more ParamKeys:

type Getter<K extends StateKey> = (state: State, stateKey: K) =>
    Partial<Record<ParamKey, ParamValue>>;
type GetterMap<KS extends StateKey> = { [K in KS]: Getter<K> };

const GETTER_MAP: GetterMap<StateKey> = {
    [StateKey.A]: (s, k) => ({ [ParamKey.A]: s[k] }),
    [StateKey.B]: (s, k) => ({ [ParamKey.B1]: s[k] }),
    [StateKey.C]: (s, k) => ({
        [ParamKey.C1]: s[k].key1,
        [ParamKey.C2]: s[k].key2,
    }),
} as const;

And combine all getters:

const getParams = <K extends StateKey>(state: State, getters: GetterMap<K>) => {
    const result: Partial<Record<ParamKey, ParamValue>> = {};
    for (const stateKey in getters) {
        Object.assign(result, getters[stateKey](state, stateKey));
    }
    return result;
};

Goal

I’d like params to have the precise inferred type::

type ExpectedParamsType = {
  a: number;      // from State.a
  b1: string[];   //    - State.b
  c1: string;     //    - State.c.key1
  c2: number;     //    - State.c.key2
};
const params: ExpectedParamsType = getParams(state, GETTER_MAP);

Instead, TypeScript currently infers it as just Partial<Record<ParamKey, ParamValue>>.

Question Is it possible, within the current implementation, to define a correct type that adapts to getters: GetterMap<K> (where we pass in GETTER_MAP) and infers the actual output type?

Or are there simpler / alternative patterns for achieving strongly typed param maps like this?

TS Playground

First, if you annotate the type of GETTER_MAP with GetterMap<StateKey>, then that's the only thing TypeScript will know about the type. Meaning all the specifics about the outputs of the getters will be lost. Instead of annotating, you mostly just want to check that GETTER_MAP is compatible with your desired type, and maybe use the desired type to contextually infer the s and k callback parameters. You can do that most easily via the satisfies operator, so instead of const vbl: Type = ⋯ you can write const vbl = ⋯ satisfies Type.

Here's how I would do that:

type GetterOutput = Partial<Record<ParamKey, ParamValue>>
type GetterMap = { [K in StateKey]: (state: State, stateKey: K) => GetterOutput };

const GETTER_MAP = {
    [StateKey.A]: (s, k) => ({ [ParamKey.A]: s[k] }),
    [StateKey.B]: (s, k) => ({ [ParamKey.B1]: s[k] }),
    [StateKey.C]: (s, k) => ({
        [ParamKey.C1]: s[k].key1,
        [ParamKey.C2]: s[k].key2,
    }),
} satisfies GetterMap;

/* const GETTER_MAP: {
    a: (s: State, k: "a") => {
        a: number;
    };
    b: (s: State, k: "b") => {
        b1: string[];
    };
    c: (s: State, k: "c") => {
        c1: string;
        c2: number;
    };
} */

First I've given a GetterOutput name to Partial<Record<ParamKey, ParamValue>> so it's less painful to refer to it. Then GetterMap doesn't have to be generic anymore since apparently you always want to use the full set of StateKey keys and not just some of them. But other than that it's very similar to what you've done. The only real difference is satisfies, and you can see that the type of GETTER_MAP contains all the information needed for getParams() to assemble the output type you care about.

Well, I also removed the const assertion, which doesn't seem to be there for any particular reason (unless you really care about the readonly modifier on your methods). Keep it if you want.

---

Now we have to give getParams() a suitable call signature. Your current one looks like

const getParams: <K extends StateKey>(
    state: State, 
    getters: Pick<GetterMap, K>
) => GetterOutput

but that's not what you want. You certainly don't care about restricting K to some subset of StateKey (you actively want to error if a key from StateKey is left out, right?), so that generic isn't doing you any good. You might as well write

const getParams: (
    state: State, 
    getters: GetterMap
) => GetterOutput

which is true but not helpful, since you want a specific subtype of GetterOutput corresponding to GetterMap. You need getParams to be generic in either the type of getters or some type related to it. My inclination is

const getParams: <T extends { [K in StateKey | keyof T]: GetterOutput; }>(
    state: State, 
    getters: { [K in keyof T]: (state: State, stateKey: K) => T[K]; }
) => MergeProperties<T>

Here getParams is generic in T, a map of all the output types of the getters in the getters map. For GETTER_MAP that will look like {a: {a: number}, b: {b1: string[]}, c: {c1: string, c2: number}}. Note that we need to constrain T to some object type where all the keys from StateKey are present, and all the properties (even ones which might not be in StateKey) are known to be assignable to GetterOutput. That looks like the recursive constraint {[K in StateKey | keyof T]: GetterOutput }. Note that while I doubt you want to support more keys than those in StateKey, it would be more complicated to prevent them. You want K in StateKey to require all the keys from StateKey, but also K in keyof T so TypeScript knows that every property is an appropriate getter output.

Then the type of getters is a mapped type over T, which maps each getter output to a getter which produces that output. And importantly, we make sure that the second parameter to the getter is the same type as the key K.

And then there's the return type, MergeProperties<T>. Which we need to define. The idea is that T is an object type where we don't care about the keys anymore, and we want to take all the values and merge them together, like with an intersection. So MergeProperties<T> should be equivalent to the intersection of all the properties of T. And for that we can use a technique similar to the one in Transform union type to intersection type :

type MergeProperties<T> = { [K in keyof T]: (x: T[K]) => void } extends
    Record<string, (x: infer I) => void> ? { [K in keyof I]: I[K] } : never

Here we're walking through each property and putting it in a contravariant position, and then inferring a single type for the combination of all the properties. This produces an intersection instead of a union. It's explained in the release notes for conditional types. So in the above, I will end up being the intersection of all the T[K]s. And then for good measure I do a simple identity mapped type over it {[K in keyof I]: I[K]} so the result will be a single object instead of an equivalent intersection of a bunch of objects (e.g., {x: 0, y: 1, z: 2} instead of {x: 0} & {y: 1} & {z: 2}.

---

Let's test it out:

const params = getParams(state, GETTER_MAP);
/* const params: {
    a: number;
    b1: string[];
    c1: string;
    c2: number;
} */ 

Looks good. Your output type is exactly what you expected. Let's make sure that getParams() works for other test cases:

getParams(state, { a: (s, k) => ({}), b: (s, k) => ({}) }); // error!
// c is missing

getParams(state, { a: (s, k) => ({}), b: (s, k) => ({}), c: (s, k) => 1 }); // error!
// Type 'number' has no properties in common with type 'Partial<Record<ParamKey, ParamValue>>'

const p1 = getParams(state, { ...GETTER_MAP, w: (s, k) => ({}) }); // okay

const p2 = getParams(state, {
    a: (s, k) => ({ b1: s[k] }),
    b: (s, k) => ({ b1: s[k] }),
    c: (s, k) => ({})
}); // okay
/* const p2: {
    b1: number & string[];
} */

This is all behaving as I intended but it might not be what you want. I consider it out of scope to deal with weird edge cases. Note that getParams() requires all the StateKey keys to be there, and will get angry if you leave one out or make it return something that's not part of GetterOutput. Extra keys are allowed. And nothing stops you from having multiple getters return the same param key with different types, and your output type will end up a weird intersection. Again, maybe this isn't perfect, and you can likely make things more restrictive, but this is the general approach I would use.

Playground link to code