How can a TypeScript forEach, infer the types of a union, and change the return type

Question

We want to take a document type and mutate some of it's fields and then return the new document with the new types. We want to determine the fields to mutate based on an array of values (collectionDataFields in the example below).

type AllDocuments = {
    __typename: 'cat'
    a: '123',
    b: '123',
    z: 'xyz'
} | {
    __typename: 'dog'
    a: '123',
    c: '123',
    z: 'xyz'
}

const collectionDateFields = {
    cat: ['a', 'b'],
    dog: ['a', 'c']
} as const

export const useStringsToNums = (document: AllDocuments) => {
  const dateFields = collectionDateFields[document.__typename]

  dateFields.forEach((field) => {
    document[field] = parseInt(document[field])
  })
  return document
}

In the example, dateFields is typed as readonly ["a", "b"] | readonly ["a", "c"] but the individual field comes in as any.

The document is returned as the AllDocuments even though we want it to have the field with the changed type (as a number).

TS Playground Example

Edit: For an example of the usage, I would do something like this:

useStringsToNums({
    __typename: 'cat'
    a: '123',
    b: '123',
    z: 'xyz'
})

In this case, the function turns the properties that match __typename: 'cat' into numbers (with the parseInt).

Answer 1

We want to take a document type and mutate some of its fields

If you have an object of a type where there is a string-valued property at some key, you can't assign a number to that same property without the compiler complaining. One of the motivations behind TypeScript is to add a strong static type system to JavaScript, and assigning a number to something which is supposed to be a string goes against that:

function doSomething(doc: AllDocuments) {
  doc.a = 123; // error! Type '123' is not assignable to type '"123"'    
}

You could just use type assertions to suppress these errors:

function doSomething(doc: AllDocuments) {
  doc.a = 123; // error! Type '123' is not assignable to type '"123"'    
}

But all this does is make things worse in the long run. TypeScript does not model arbitrary type mutations, only narrowing. There's no way to say that after calling doSomething(doc), the type of doc has changed so that a former string property is now a number. By asserting that a number is a string, you are lying to the compiler, and you are liable to encounter errors at runtime:

const doc: AllDocuments = { __typename: "cat", a: "123", b: "123" };
doSomething(doc);
try {
  doc.a.toUpperCase(); // compiles but has runtime error:
} catch (e) {
  console.log(e); // doc.a.toUpperCase is not a function    
}

Assuming you don't want to widen those property types from "123" to something that can handle both "123" and 123, then it's probably a better idea to just create a new object of the type you're looking for, and treat the original object as immutable:

function doSomethingElse(doc: AllDocuments) {
  return { ...doc, a: 123 }
}
const newDoc = doSomethingElse(doc);
/* const newDoc: {
  a: number;
  __typename: 'cat';
  b: '123';
 } | {
  a: number;
  __typename: 'dog';
  c: '123';
 } */
console.log(newDoc.a.toFixed(2)); // "123.00"

For the rest of this answer I will use this approach.

---

---

Now for the rest of your question. You are running into a problem with what I've been calling correlated union types as discussed in microsoft/TypeScript#30581.

In the following code:

  const dateFields = collectionDateFields[document.__typename]
  dateFields.forEach((field: "a" | "b" | "c") => {
    document[field] // oops
  })

even if we annotate field with the type we expect it to be, the compiler is unable to see that document has a key named field. The type of document is a union of two distinct types, and the type of dateField is also a union of two distinct types, and the type of field is also a union. But the compiler sees these three unions as independent or uncorrelated. As far as it understands, document.__typename could be "cat" and yet field could be "c". You know that this is impossible, but the compiler loses the thread.

---

Sometimes it's possible to refactor code to avoid the correlated-union problem. Unfortunately, despite spending some time trying to get something sort of type safe from the compiler, the closest I could get isn't useful enough to be worth the complexity:

const useDatesToLuxon = <T extends AllDocuments['__typename']>(
  document: Extract<AllDocuments, { __typename: T }>
) => {
  type Fields = (typeof collectionDateFields)[T][number];
  const dateFields: readonly Fields[] = collectionDateFields[document.__typename]
  const doc = document as Record<Fields, string>;
  const newDoc = {} as Record<Fields, number>;
  dateFields.forEach((field) => {
    newDoc[field] = parseInt(doc[field])
  })
  const ret = { ...(document as Omit<typeof document, Fields>), ...newDoc };
  type Ret = typeof ret;
  return ret as { [K in keyof Ret]: Ret[K] }
}

That's a lot of type juggling, and in the end it only works for arguments known to be a single member of the AllDocuments union:

const r = useDatesToLuxon({ __typename: "cat", a: "123", b: "123" });
/* const r: {
    __typename: 'cat';
    a: number;
    b: number;
} */
console.log(r); // as expected

If you try to call it on an object whose type is only known to be the full AllDocuments union, you get the wrong output type:

function butWait(doc: AllDocuments) {
  const u = useDatesToLuxon(doc);
  /* const u: {
    __typename: "cat" | "dog";
    a: number;
    b: number;
    c: number;
    } */
  // UGH, no, that's not the type we want; we cannot programmatically *distribute* the 
  // behavior of useDatesToLuxon across the union members of AllDocuments
}

So let's give up on such refactoring.

---

There isn't a great solution to correlated unions. One way you can work around the problem is by using redundant code, breaking your code block into several identical chunks where the compiler can use control flow analysis to go through the different cases:

export const useDatesToLuxon = (document: AllDocuments) => {
  switch (document.__typename) {
    case "cat": {
      const dateFields = collectionDateFields[document.__typename]
      const temp = {} as Record<typeof dateFields[number], number>
      dateFields.forEach((field) => {
        temp[field] = parseInt(document[field])
      })
      return { ...document, ...temp }
    }
    case "dog": {
      const dateFields = collectionDateFields[document.__typename]
      const temp = {} as Record<typeof dateFields[number], number>
      dateFields.forEach((field) => {
        temp[field] = parseInt(document[field])
      })
      return { ...document, ...temp }
    }
  }
}

This works very well, and is quite type safe. In the "cat" case and in the "dog" case, the compiler knows exactly what is happening. But such redundancy doesn't scale. It would be nice if you could ask the compiler to pretend you had written such redundant code, as I asked for in microsoft/TypeScript#25051, but you can't.

---

My general recommendation for this situation is to use type assertions to loosen the type checking inside the implementation of your function. This shifts the burden of maintaining type safety away from the compiler (which can't do it properly here) onto you, so you need to be careful.

Here's one possible way to do it. First, let's compute the type we expect to come out of your function:

type CollectionDateFields = typeof collectionDateFields;

type ConvertedDocuments = { 
  [D in AllDocuments as D['__typename']]: { 
    [K in keyof D]: K extends 
      CollectionDateFields[D['__typename']][number] ? number : D[K] 
} }[keyof CollectionDateFields]

That's using key remapping to iterate over each member of AllDocuments and converting the relevant properties mentioned in collectionDateFields to number. You can verify that it's correct:

/* type ConvertedDocuments = {
__typename: "cat";
a: number;
b: number;
} | {
__typename: "dog";
a: number;
c: number;
} */

Now we can give useDatesToLuxon a single overload call signature describing what it does:

function useDatesToLuxon<D extends AllDocuments>(
  document: D
): Extract<ConvertedDocuments, Pick<D, '__typename'>>

That's a generic function which can turn the whole AllDocuments union into the whole ConvertedDocuments union, as well as turning any subtype into the analogous output type. Now for the implementation:

function useDatesToLuxon(_document: AllDocuments): ConvertedDocuments {
  // let's just pretend it's one of the members of the union
  const document = _document as Extract<AllDocuments, { __typename: "cat" }>
  const dateFields = collectionDateFields[document.__typename]
  const temp = {} as Record<typeof dateFields[number], number>
  dateFields.forEach((field) => {
    temp[field] = parseInt(document[field])
  })
  return { ...document, ...temp }
}

Here the easiest assertion for me to make is to just pretend that document is one of the members of the union. It's not technically true, but it suppresses the complaints from the compiler. We just have to be careful that we've implemented it properly.

And now to test it:

const r = useDatesToLuxon({ __typename: "cat", a: "123", b: "123" });
/* const r: {
    __typename: 'cat';
    a: number;
    b: number;
} */
console.log(r); // as expected

function andThen(doc: AllDocuments) {
  const u = useDatesToLuxon(doc);
  /* {
    __typename: "cat";
    a: number;
    b: number;
  } | {
    __typename: "dog";
    a: number;
    c: number;
  } */

}

Looks good!

Playground link to code