Type safe function that receives an object and a string and returns a new object with string as a ke-CodePudding

I have a function that adds a field to an object and returns the object. I want to type the function so that when I use the function I'd have an object that has original properties and the added one.

So far I have:

function addComputedField<T extends Record<string, unknown>, V>(
  document: T,
  fieldName: string,
  computation: () => V
): T & { [field: typeof fieldName]: V };
function addComputedField<T extends Record<string, unknown>, V>(
  document: T[],
  fieldName: string,
  computation: (doc: T) => V
): (T & { [field: typeof fieldName]: V })[];

function addComputedField<T extends Record<string, unknown>, V>(
  document: T | T[],
  fieldName: string,
  computation: (doc?: T) => V
) {
  let result: (T & { [field: typeof fieldName]: V }) | (T & { [field: typeof fieldName]: V })[];

  if (Array.isArray(document)) {
    result = document.map((doc) => ({
      ...doc,
      [fieldName]: computation(doc),
    }));
  } else {
    result = { ...document, [fieldName]: computation() };
  }

  return result;
}

No errors, good, but it doesn't work as expected.

For example, if I use it with these arguments:

const obj = { a: 1, b: 2, c: 3 }
const withComputed = addComputedField(obj, "d", () => 4)

type of withComputed would be

const withComputed: {
    a: number;
    b: number;
    c: number;
} & {
    [field: string]: number;
}

instead of expected type

const withComputed: {
    a: number;
    b: number;
    c: number;
} & {
    d: number;
}
// or 
const withComputed: {
    a: number;
    b: number;
    c: number;
    d: number;
}

How can I achieve expected results?

CodePudding user response：

The reason this doesn't work is because typeof fieldName is just string - it isn't inferred as something stricter than that based on the argument you call the function with.

The behaviour you want requires an additional generic type parameter for the field name:

function addComputedField<T extends object, K extends PropertyKey, V>(
  document: T,
  fieldName: K,
  computation: () => V
): T & Record<K, V>;
// ...