Type-safe generic reduce util to create object from array

Question

I'm looking for a generic and type-safe way to model the following JavaScript in TypeScript:

const records = [
  { name: "foo", id: 1, data: ["foo"] },
  { name: "bar", id: 2, data: ["bar"] },
  { name: "baz", id: 3, data: ["baz"] }
];

function keyBy(collection, k1, k2) {
  if (k2) {
    return collection.reduce((acc, curr) =>
      ({ ...acc, [curr[k1]]: curr[k2] }), {});
  } else {
    return collection.reduce((acc, curr) =>
      ({ ...acc, [curr[k1]]: curr }), {});
  }
}

console.log(keyBy(records, "name", "data"));
// { foo: [ 'foo' ], bar: [ 'bar' ], baz: [ 'baz' ] }

console.log(keyBy(records, "name"));
// {
//   foo: { name: 'foo', id: 1, data: [ 'foo' ] },
//   bar: { name: 'bar', id: 2, data: [ 'bar' ] },
//   baz: { name: 'baz', id: 3, data: [ 'baz' ] }
// }

The idea is to create a util that will reduce an array into an object keyed by the value at a given key, and with a value of either the entire object, or optionally a specific data point at a given second key (this explanation may be a bit poor, but hopefully the example speaks for itself).

This is pretty simple JS, but seems hard to get the types right in TS. Here's what I've come up with so far, but I've needed to create two functions in order to get the return types right and if all feels a bit hacky. I was unable to get a conditional return type to work here, so am OK with two functions if that's the way it has to be, but wondering if there's a better approach here (perhaps something that could result in Record<T[K], T> or Record<T[K], T[K2]> rather than the record being keyed by ObjectKey). Thanks.

type ObjectKey = string | number | symbol;

const isValidKey = (x: any): x is ObjectKey =>
  typeof x === "string" || typeof x === "number" || typeof x === "symbol";

function keyBy<T extends object, K extends keyof T>(collection: T[], key: K) {
  return collection.reduce((acc, curr) => {
    const valueAtKey = curr[key];

    if (isValidKey(valueAtKey)) {
      return { ...acc, [valueAtKey]: curr };
    }

    throw new Error("T[K] is not a valid object key type");
  }, {} as Record<KeyType, T>);
}

function keyByWith<T extends object, K extends keyof T, K2 extends keyof T>(
  collection: T[],
  k: K,
  k2: K2,
) {
  return collection.reduce((acc, curr) => {
    const valueAtKey = curr[k];

    if (isValidKey(valueAtKey)) {
      return { ...acc, [valueAtKey]: curr[k2] };
    }

    throw new Error("T[K] is not a valid object key type");
  }, {} as Record<ObjectKey, T[K2]>);
}

P.S. I know lodash has a similar keyBy function, but I don't think they have anything similar to keyByWith shown above.

Answer 1

The biggest problem is that records is being inferred as type:

{
    name: string;
    id: number;
    data: string[];
}[]

Which means keyBy(records, 'name') can only give you back string. If you add a as const assertion to records, then you can get some literal strings and you have stronger types to work with.

const records = [
  { name: "foo", id: 1, data: ["foo"] },
  { name: "bar", id: 2, data: ["bar"] },
  { name: "baz", id: 3, data: ["baz"] }
] as const;

---

Then you need to type your reduce'd result object as

Record<T[K] & ObjectKey, T>

or

Record<T[K] & ObjectKey, T[K2]>

so that the keys from the generic T are used.

The T[K] & ObjectKey with an invalid key type will resolve to never, but you will also throw a runtime exception there so that doesn't matter much.

---

And lastly, you can can use overloading to declare multiple signatures to make this one function. This will have two signatures:

// One key
function keyBy<
  T extends object,
  K extends keyof T
>(
  collection: readonly T[],
  key: K
): Record<T[K] & ObjectKey, T>

// Two keys
function keyBy<
  T extends object,
  K extends keyof T,
  K2 extends keyof T
>(
  collection: readonly T[],
  k: K,
  k2: K2,
): Record<T[K] & ObjectKey, T[K2]>

And an implementation with something like:

// Implementation
function keyBy<
  T extends object,
  K extends keyof T,
  K2 extends keyof T
>(
  collection: readonly T[],
  k: K,
  k2?: K2,
): Record<T[K] & ObjectKey, T[K2]> | Record<T[K] & ObjectKey, T> {
  return collection.reduce((acc, curr) => {
    const valueAtKey = curr[k];

    if (isValidKey(valueAtKey)) {
      if (k2) return { ...acc, [valueAtKey]: curr[k2] };
      return { ...acc, [valueAtKey]: curr };
    }

    throw new Error("T[K] is not a valid object key type");
  }, {} as Record<T[K] & ObjectKey, T[K2]> | Record<T[K] & ObjectKey, T>);
}

And now this works:

const testA = keyBy(records, "name");
testA.foo.data // readonly ["foo"] | readonly ["bar"] | readonly ["baz"]

const testB = keyBy(records, "name", "data");
testB.foo // readonly ["foo"] | readonly ["bar"] | readonly ["baz"]

Playground

Answer 2

Building off Alex's answer it is actually possible to infer this type fully and discriminate it properly, using mapped object types. But it definitely is more verbose and requires some massaging.

const testA = keyBy(records, "name");
testA.foo.data // readonly ["foo"]

const testB = keyBy(records, "name", "data");
testB.foo // readonly ["foo"]

I went ahead and took some tools from other answers to achieve this

//https://stackoverflow.com/questions/61410242/is-it-possible-to-exclude-an-empty-object-from-a-union
type AtLeastOne<T, U = {[K in keyof T]: Pick<T, K> }> = Partial<T> & U[keyof U];

type ExcludeEmpty<T> = T extends AtLeastOne<T> ? T : never; 

function keyBy<
  T extends Record<any, any>,
  K extends keyof T,
  K2 extends keyof T
>(
  collection: readonly T[],
  k: K,
): {
  [P in T[K]]: ExcludeEmpty<{
    [P2 in keyof T as T[K] extends P ? T[K] : never]: T
  }>[P]
}
function keyBy<
  T extends Record<any, any>,
  K extends keyof T,
  K2 extends keyof T
>(
  collection: readonly T[],
  k: K,
  k2: K2,
): {
  [P in T[K]]: ExcludeEmpty<{
    [P2 in keyof T as T[K] extends P ? T[K] : never]: T[K2]
  }>[P] 
}
// Implementation
function keyBy<T extends Record<any, any>, K extends keyof T, K2 extends keyof T>(
  collection: readonly T[],
  k: K,
  k2?: K2,
): {
  [P in T[K]]: ExcludeEmpty<{
    [P2 in keyof T as T[K] extends P ? T[K] : never]: T
  }>[P]
} | {
  [P in T[K]]: ExcludeEmpty<{
    [P2 in keyof T as T[K] extends P ? T[K] : never]: T[K2]
  }>[P] 
} {...}

View this on TS Playground

Answer 3

This isn't exactly an answer to my question since it's a completely different approach, but I've been playing around with this a bit more and thought it worth sharing an alternative solution to the same problem that uses callbacks, not keys, to extract both the key and value of the resulting object.

This approach has the same level of type safety as @AlexWayne's accepted answer, though not as much as @CodyDuong's.

However, it supports greater flexibility in terms of data transformation for the object's values (rather than being limited to T or T[K]), and does not require a runtime check to ensure the key of the object is a valid object key (rather it will just fail compilation if an invalid key key extractor is provided):

type User = {
  id: number;
  name: string;
  data: string[] | undefined;
};

const records: User[] = [
  { id: 1, name: "foo", data: ["fee"] },
  { id: 2, name: "baz", data: ["fi"] },
  { id: 3, name: "bar", data: undefined },
];

type ObjectKey = string | number | symbol;

type ToKey<T extends object, U extends ObjectKey> = (
  value: T,
  index: number
) => U;

type ToValue<T extends object, V> = (value: T, index: number, arr: T[]) => V;

function keyBy<T extends object, K extends ObjectKey>(
  collection: T[],
  keyExtractor: ToKey<T, K>
): Record<K, T>;

function keyBy<T extends object, K extends ObjectKey, V>(
  collection: T[],
  keyExtractor: ToKey<T, K>,
  valueExtractor?: ToValue<T, V>
): Record<K, V>;

function keyBy<T extends object, K extends ObjectKey, V>(
  collection: T[],
  keyExtractor: ToKey<T, K>,
  valueExtractor?: ToValue<T, V>
) {
  return collection.reduce<Record<K, T> | Record<K, V>>(
    (acc, curr, index, arr) => ({
      ...acc,
      [keyExtractor(curr, index)]: valueExtractor
        ? valueExtractor(curr, index, arr)
        : curr,
    }),
    {} as any
  );
}

const nameToData = keyBy(
  records,
  (x) => x.name,
  (x) => x.data
);

const nameToUser = keyBy(records, (x) => x.name);

const indexToUser = keyBy(records, (x, i) => i);

const indexToName = keyBy(
  records,
  (x, i) => i,
  (x) => x.name
);

const idToTransformedData = keyBy(
  records,
  (x, i) => i,
  (x) => x.data?.map((s) => [s.repeat(3)])
);

TS Playground