Recursively transform all leafs of an object tree in typescript

Question

Given a simple object tree that either contains values of its own type or of a type that needs to be transformed:

interface Tree<Leaf> {
  [key: string]: Tree<Leaf> | Leaf;
}

I want to define a function that recursively transforms all leafs to another type like this:

function transformTree<S, T>(
  obj: Tree<S>,
  transform: (value: S) => T,
  isLeaf: (value: S | Tree<S>) => boolean
): Tree<T> {
  return Object.assign(
    {},
    ...Object.entries(obj).map(([key, value]) => ({
      [key]: isLeaf(value)
        ? transform(value as S)
        : transformTree(value as Tree<S>, transform, isLeaf),
    }))
  );
}

How do can I maintain the types of the leafs between the source tree and transformed tree?

Testing the above doesn't work:

class Wrapper<T> {
  constructor(public value: T) {}
}

function transform<T>(wrapped: Wrapper<T>): T {
  return wrapped.value;
}

function unwrap<T>(wrapped: Tree<Wrapper<T>>): Tree<T> {
  return transformTree<Wrapper<T>, T>(
    wrapped,
    transform,
    (value: Wrapper<T> | Tree<Wrapper<T>>) => value instanceof Wrapper
  );
}

const obj = unwrap<string>({
  foo: {
    bar: new Wrapper("baz"),
  },
  cow: new Wrapper("moo"),
});

function handleBaz(value: "baz") {
  return true;
}

function handleMoo(value: "moo") {
  return true;
}

handleBaz(obj.foo.bar); // Error:(162, 19) TS2339: Property 'bar' does not exist on type 'string | Tree<string>'. Property 'bar' does not exist on type 'string'.
handleMoo(obj.cow); //Error:(163, 11) TS2345: Argument of type 'string | Tree<string>' is not assignable to parameter of type '"moo"'. Type 'string' is not assignable to type '"moo"'.

I can see that the error occurs, because the tree structure is not maintained through the transformation (the transformation only works at runtime). But given that the transformation is predictable given a known structure of the input tree, I feel that there should be a way to do this in typescript.

Here is my attempt to solve this:

type TransformTree<Leaf, InputTree extends Tree<Leaf>, T> = {
  [K in keyof InputTree]: InputTree[K] extends Leaf
    ? T
    : InputTree[K] extends Tree<Leaf>
    ? TransformTree<Leaf, InputTree[K], T>
    : never;
};

type IsLeaf<Leaf> = (value: Tree<Leaf> | Leaf) => boolean;

function transformTree<T extends Tree<From>, From, To>(
  tree: T,
  transform: (value: From) => To,
  isLeaf: IsLeaf<From>
): TransformTree<From, typeof tree, To> {
  return Object.assign(
    {},
    ...Object.entries(tree).map(([key, value]) => ({
      // XXX have to cast the value in each case, because typescript cannot predict
      // the outcome of isLeaf().
      [key]: isLeaf(value)
        ? transform(value as Extract<typeof tree[typeof key], From>)
        : transformTree(
            value as Extract<typeof tree[typeof key], Tree<From>>,
            transform,
            isLeaf
          ),
    }))
  );
}

It still doesn't seem to understand the nested types:

function unwrap<T>(
  wrapped: Tree<Wrapper<T>>
): TransformTree<Wrapper<T>, typeof wrapped, T> {
  return transformTree(
    wrapped,
    transform,
    (value: Wrapper<T> | Tree<Wrapper<T>>) => value instanceof Wrapper
  );
}


function handleBaz(value: "baz") {
  return true;
}

function handleMoo(value: "moo") {
  return true;
}

handleMoo(obj.cow); // OK

handleBaz(obj.foo.bar); // TS2339: Property 'bar' does not exist on type 'never'.

Playground Link

It seems like typescript still thinks that if the field value is not a Leaf, then it may be something other than a sub-tree.

Answer 1

In what follows I'm only going to worry about typings and not implementation. All functions will just be declared, as if the actual implementations are in some JS library and these are their declaration files.

Also, your isLeaf functions should probably be typed as user-defined type guards whose return type is a type predicate instead of just boolean. A function signature isLeaf: (value: S | Tree<S>) => value is S is similar to one that returns boolean, except that the compiler will actually understand that in if (isLeaf(x)) { x } else { x }, the x in the true block will be S and the x in the false block will be Tree<S>.

Okay, here goes:

---

The type Tree<X> is too general to keep track of specific key and value types. All the compiler knows about a value of the type, say, Tree<string>, is that its is an object type whose properties are either of type string or Tree<string>. Once you do, say, this:

const x: Tree<string> = { a: "", b: { c: "", d: { e: "" } } };

You've thrown away all the details about the particular stucture as well as all the details about any specific subtypes of string at the leaves:

x.a.toUpperCase(); // error
x.z; // no error 

If all you cared about were to come up with a type transformation that maintained the nested key structure and turned some subtype of Tree<X> into a subtype of Tree<Y> with the same shape, you could do it. But in the most straightforward implementation, all the leaves of the resulting tree will be of type Y and not any narrower type. Here's how I'd write it:

type TransformTree<T extends Tree<X>, X, Y> = { [K in keyof T]:
    T[K] extends X ? Y :
    T[K] extends Tree<X> ? TransformTree<T[K], X, Y> :
    T[K];
};
declare function transformTree<X, Y, TX extends Tree<X>>(
    obj: TX,
    transform: (value: X) => Y,
    isLeaf: (value: X | Tree<X>) => value is X
): TransformTree<TX, X, Y>;

And you can see it work on a simple example like this:

const t1 = { a: "A", b: { c: "CC", d: { e: "EEE" } } };
const t2 = transformTree(t1,
    (x: string) => x.length,
    (v): v is string => typeof v === "string"
); 
t2.a; // number
t2.b.c; // number
t2.b.d.e; // number

---

But you want something significantly more ambitious here; it seems like you want to have the leaf transformation map not just from a specific type X to a specific type Y, but you'd like to specify some general type function like type F<T extends X> = ... and map the leaf from type Z extends X to F<Z>.

In your example, your input type is something like Wrapped<any>, and your output type function would look like type F<T extends Wrapped<any> = T["value"];

Unfortunately, this more general type of transformation cannot be expressed in TypeScript. You can't have a type function like type TransformType<T extends Tree<X>, X, F> = ... where F is itself a type function that takes a parameter. That would require the language to support what's known as "higher kinded types". There is a longstanding open feature request for this at microsoft/TypeScript#1213, and while it would be amazing to have them, it doesn't look like it will happen in the foreseeable future.

---

What you can do is envision particular types of leaf type transformations and implement specific versions of TransformTree for them. For example, if your leaf type mapping just indexes into a single property like type F<T extends Record<K, any>, K extends PropertyKey> = T[K], as in your Unwrap case, then you can write it like this:

type TransformTreeIdx<T, X, K extends keyof X> = { [P in keyof T]:
    T[P] extends X ? X[K] :
    TransformTreeIdx<T[P], X, K>;
};
declare function transformTreeIdx<TX, X, K extends keyof X>(
    obj: TX,
    key: K,
    isLeaf: (value: any) => value is X
): TransformTreeIdx<TX, X, K>;

And then use it:

const w = {
    foo: {
        bar: new Wrapper("baz" as const),
    },
    cow: new Wrapper("moo" as const),
};

const w2 = transformTreeIdx(
    w, "value", (x: any): x is Wrapper<any> => x instanceof Wrapper
);

handleBaz(w2.foo.bar);
handleMoo(w2.cow);

---

Or, as mentioned in your comments, you might want to do the opposite with a particular generic interface/class, like say, Wrapper... the mapping is to convert Z extends X into Wrapper<Z>:

type TransformTreeWrap<T, X> = { [P in keyof T]:
    T[P] extends X ? Wrapper<T[P]> :
    TransformTreeWrap<T[P], X>;
};
declare function transformTreeWrap<TX, X, K extends keyof X>(
    obj: TX,
    isLeaf: (value: any) => value is X
): TransformTreeWrap<TX, X>;

and use it:

const u = {
    foo: {
        bar: "baz" as const,
    },
    cow: "moo" as const,
};

const u2 = transformTreeWrap(
    u, (x: any): x is string => typeof x === "string"
);
handleBaz(u2.foo.bar.value);
handleMoo(u2.cow.value);

---

Or, you might have an array of isLeaf/transform pairs that allow each node to be tested against different more specific transformations. So, for example, any time you find a "moo" value in the tree you output a number and any time you find a "baz" value in the tree you output a boolean.

Then your typings could look like:

type TransformTreeMap<T, M extends [any, any]> = { [K in keyof T]:
    T[K] extends M[0] ? Extract<M, [T[K], any]>[1] :
    TransformTreeMap<T[K], M> };

type IsLeafAndTransformer<I, O> = {
    isLeaf: (x: any) => x is I,
    transform: (x: I) => O
}
type TransformArrayToMap<M extends Array<IsLeafAndTransformer<any, any>>> = {
    [K in keyof M]: M[K] extends IsLeafAndTransformer<infer I, infer O> ?
    [I, O] : never }[number]

declare function transformTreeMap<T, M extends Array<IsLeafAndTransformer<any, any>>>(
    obj: T,
    ...transformers: M
): TransformTreeMap<T, TransformArrayToMap<M>>;

and you use it:

const mm = transformTreeMap(u,
    { isLeaf: (x: any): x is "moo" => x === "moo", transform: (x: "moo") => 123 },
    { isLeaf: (x: any): x is "baz" => x === "baz", transform: (x: "baz") => true }
);
mm.cow // number
mm.foo.bar // boolean

---

So there are really a lot of different and even quite powerful tree transformations you can define... just not the fully generic higher-order type implied by the question. Hopefully one of those will give you a way forward. Okay, good luck!

Playground link to code

Answer 2

I have a function that that takes an object and returns a clone with all leafs as strings:

function leafsToString<O extends Record<any, any>>(obj: O): unknown {
    const newObj = obj instanceof Array ? [] : Object.create(null);
    for (const key of Object.keys(obj)) {
        if (typeof obj[key] === 'object') {
            newObj[key] = leafsToString(obj[key]);
        } else {
            newObj[key] = obj[key] + '';
        }
    }
    return newObj;
}

I just need a way to type the result. So I came up with this:

type LeafsToString<O> = {
    [K in keyof O]:
        O[K] extends Record<any, any>
        ? LeafsToString<O[K]>
        : string
}

Playground link with examples.