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

Just use a type assertion, your concern should be the call site, that is typed correctly and gives you errors where it should. The implementation can't really be typed correctly if you want to assign a specific value inside the function.

You can make the result of T[K] extend number, for example but adding a constraint to T of Record<K, number>, but we still not be able to assign concrete values to obj[key]

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

function f<T extends Record<K, number>, K extends KeyOfType<T, number>>(obj: T, key: K, value: T[K]) {
    let r = obj[key]; 
    r.toExponential(); // seems numberish, but it `T[K]` which does extend number, but might not be number
    obj[key] = obj[key] // T[K] is assignable to T[K]
    obj[key] = value; // even if it is a parameter 
    obj[key] = 5; // still an error
}

declare const foo: Foo;
f(foo, "a", 1); // Allowed, good
f(foo, "b", 2); // Error, good

const other: { 
  a: 1
} = {
  a: 1
}
f(other, "a", 1) // this will break the type of other because of obj[key] = 5

Playground Link

The reason this is so, is the last example f(other, "a", 1). Here a in other has type 1, which does extend number, so f(other, "a", 1) is a valid call to f, but inside we want to assign other[key] = 5. This would break the type of other. The problem here is that there is no way so specify that T[K] has an upper constraint of number just a lower constraint.

Answer 2

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.