Creating a fluent, stateful builder using the TypeScript type system

Question

I would like to create a builder which can represent builder.a().b().build() or builder.b().a().build(), but not builder.a().build(), builder.b().build() or builder.a().a().build(), etc.

Obviously I can do the validation in the build method, but I would like the compiler to hint at this (and for vs code to provide auto complete). I think the TS type system can represent this using mapped types, unions and intersections, but I can't quite figure it out.

Does anyone know how I could do this?

Answer 1

Let's start with the easiest part, the actual implementation. The class uses this as return type. The actual magic will happen using an additional type later on:

class ConcreteBuilder {
  a(): this {
    return this;
  }

  b(): this {
    return this;
  }

  build(): string {
    return 'foo';
  }
}

Next we'll need a generic type where we can pass in a function type and it will exchange the return type for something else. You'll see why in a minute. I copied this code from another answer:

type ArgumentTypes<T> = T extends (... args: infer U) => infer R ? U: never;
type ReplaceReturnType<O, N> = (...a: ArgumentTypes<O>) => N;

Now it gets interesting, here's the Builder type I came up with:

type Builder<K extends keyof ConcreteBuilder> = {
  [U in K]: U extends 'build' ? ConcreteBuilder[U] : ReplaceReturnType<
    ConcreteBuilder[U],
    Builder<Exclude<K, U> extends never ? 'build' : Exclude<K, U>>
  >
};

The type has a generic argument K which is restricted to being a key of ConcreteBuilder. This could be 'a' or also 'a' | 'b' | 'build'. We'll use this argument to determine which methods should be available in the resulting type.

For this we use a mapped type that iterates over all keys in K. For each U in K, except the build method we modify the method's return type.

ConcreteBuilder[U] refers to the original function type of the method.

If we've reached the build method (U extends 'build') we keep the original type from ConcreteBuilder. Otherwise we use ReplaceReturnType to preserve the original arguments but replace the return type with:

Builder<Exclude<K, U> extends never ? 'build' : Exclude<K, U>>

Lot's to unpack here. So as you can see we're replacing the return type with Builder. The arguments to which, as discussed, define what methods are available. Since this is the return type of method U, we want to remove U from the list of available methods to prevent calling it twice. This is done using the Exclude type. In our case we remove U from the union of keys K.

Because this recursive type goes on until all methods are removed from K we also need a base case. That's what the conditional type here is for. We check if the remaining keys (Exclude<K, U>) extend never which essentially means if the is "empty". And if that's the case return a builder with the method build.

Finally, the only thing left is the builder function:

function builder(): Builder<Exclude<keyof ConcreteBuilder, 'build'>> {
  return new ConcreteBuilder() as any;
}

It returns a Builder with all methods except build. There is an any cast necessary, because the type of ConcreteBuilder is less restrictive than Builder.

---

Full code

type ArgumentTypes<T> = T extends (... args: infer U) => infer R ? U: never;
type ReplaceReturnType<O, N> = (...a: ArgumentTypes<O>) => N;

type Builder<K extends keyof ConcreteBuilder> = {
  [U in K]: U extends 'build' ? ConcreteBuilder[U] : ReplaceReturnType<
    ConcreteBuilder[U],
    Builder<Exclude<K, U> extends never ? 'build' : Exclude<K, U>>
  >
};

function builder(): Builder<Exclude<keyof ConcreteBuilder, 'build'>> {
  return new ConcreteBuilder() as any;
}

class ConcreteBuilder {
  a(): this {
    return this;
  }

  b(): this {
    return this;
  }

  build(): string {
    return 'foo';
  }
}

builder().a().b().build(); // ok
builder().b().a().build(); // ok
builder().build(); // error
builder().a().build(); // error
builder().b().build(); // error
builder().a().a().build(); // error
builder().b().b().build(); // error
builder().a().b().b().build(); // error
builder().a().a().b().build(); // error

Playground