Constrain one Typescript generic parameter based on properties of another?

Question

I'm trying to write a function that takes an object and a (string) key, then operates on a property of the object. This is easy:

function f<T extends any, K extends keyof T>(obj: T, key: K) {
  const prop = obj[key]; // prop is typed as T[K]
}

I would like to constrain the key passed to the call, at compile-time, based on the type of T[K]. I tried this:

function f<T extends any, K extends keyof T>(obj: T, key: T[K] extends number ? K : never) {
  obj[key] = 5; // error, "is not assignable to" etc
}

prop is typed as T[T[K] extends number ? K : never] which reads to me like it should collapse to just number, but it does not.

My goal is to be sure that obj[key] is typed as number, inside the function, and also have calls like f({a: true}, "a") flagged as an error. Is this possible? I thought I might need to move the constraint from the function parameter declaration, to the generic parameter declaration, but I couldn't figure out the syntax.

---

ETA yet again: Playground example -- updated to try the approach suggested by @reactgular in a comment:

type AssignableKeys<T, ValueType> = {
  [Key in keyof T]-?: ValueType extends T[Key] | undefined ? Key : never
}[keyof T];

type PickAssignable<T, ValueType> = Pick<T, AssignableKeys<T, ValueType>>;

type OnlyAssignable<T, ValueType> = {
  [Key in AssignableKeys<T, ValueType>]: ValueType
};

interface Foo {
  a: number;
  b: string;
  nine: 9;
  whatevs: any;
}

type FooNumberKeys = AssignableKeys<Foo, number>; // "a" | "whatevs"
type CanAssignNumber = PickAssignable<Foo, number>; //  { a: number; whatevs: any; }
type DefinitelyJustNumbers = OnlyAssignable<Foo, number>; //  { a: number; whatevs: number; }

function f1<T>(obj: OnlyAssignable<T, number>, key: keyof OnlyAssignable<T, number>) {
  obj[key] = Math.random(); // Assignment is typed correctly, good
}

function f2<T extends object, K extends keyof PickAssignable<T, number>>(obj: T, key: K) {
  obj[key] = Math.random(); // Uh oh, Type 'number' is not assignable to type 'T[K]'.(2322)
}

declare const foo: Foo;
f1(foo, "a"); // Allowed, good
f1(foo, "whatevs"); // Allowed, good
f1(foo, "nine"); // Uh oh, should error, but doesn't!
f1(foo, "b"); // Error, good

f2(foo, "a"); // Allowed, good
f2(foo, "whatevs"); // Allowed, good
f2(foo, "nine"); // Error, good
f2(foo, "b"); // Error, good

In the Playground, DefinitelyJustNumbers shows a tooltip of {a: number; whatevs: number} -- anything that I can assign a number to is explicitly typed as number. This fixes the assignment inside the function body, but fails to detect the fact that nine is only a subset of number and so should not be allowed.

CanAssignNumber shows a tooltip of {a: number; whatevs: any}, correctly excluding nine because it's not assignable to number. This looks good, but still doesn't fix assignment inside the function f2.

Answer 1

You should extends from keys that result in a value of type number.

export type PickByValue<T, ValueType> = Pick<
  T,
  { [Key in keyof T]-?: T[Key] extends ValueType ? Key : never }[keyof T]
>;

function f<T extends object, K extends keyof PickByValue<T, number>>(obj: T, key: K) : T[K] {
  return obj[key]
}

Edit: What you're trying to do is not possible in TS AFAIK, and sometimes there's a good reason for that. let's suppose you have below code:

function f<T extends object, K extends keyof PickByValue<T, number>>(obj: T, key: K) {
    obj[key] = 5; // Type 'number' is not assignable to type 'T[K]'.
} 

const obj = {a: 9} as const;

f(obj, "a")

For example in the above scenario, value of property a is a number however it's not of type number but of type 9. there's no way for typescript to know this before hand. in other scenarios, the only thing that comes to my mind is using Type Assertions.