[SOLVED] Assigning types to a function that returns an arbitrary property of a randomly chosen element

Assigning types to a function that returns an arbitrary property of a randomly chosen element

I have this function:

function randomArrayElement<T> (anArray: T[]):T {
  return anArray[Math.floor(Math.random() * anArray.length)]
}

I'd like to add functionality where if I supply a second argument that names a property, the function will return the value of that property on the randomly chosen array element. In JS:

function randomArrayElement (anArray, elementProp = null) {
  return elementProp
    ? anArray[Math.floor(Math.random() * anArray.length)][elementProp]
    : anArray[Math.floor(Math.random() * anArray.length)]
}

How do I type that elementProp when it could be anything (or not present)? Would it be preferable to select the element's property where I call the function?

I've looked at this answer, but I want to add a step.

Solution

I think the following works as you want it to:

// Overload: no property given -> return T
function randomArrayElement<T>(arr: T[]): T;

// Overload: property key given -> return that property's type
function randomArrayElement<T, K extends keyof T>(arr: T[], prop: K): T[K];

function randomArrayElement<T, K extends keyof T>(
  arr: T[],
  prop?: K
): T | T[K] {
  const el = arr[Math.floor(Math.random() * arr.length)];
  return prop ? el[prop] : el;
}