typescripttypescript-typingstypescript-generics

Replicating the mongodb project aggregation in Typescript


My scenario: I am trying to build a typesafe wrapper around node-mongodb driver. I am having a hard time figuring out the return type for project aggregation stage.

Have a look at the TS Playground here

class BaseAggregate<Coll> {
  private pipeline: Pipeline<Coll>[] = [];
  constructor(initialPipeline?: typeof this.pipeline) {
    if (initialPipeline) {
      this.pipeline = initialPipeline;
    }
  }

  match(stageInput: DeepPartial<Coll>) {
    return new BaseAggregate<Coll>([...this.pipeline, { $match: stageInput }]);
  }

  projectActual(stageInput: Record<keyof Coll, 0 | 1>) {
    return new BaseAggregate<T>>([
      //  what goes here ----^ instead of T?
      //   ...this.pipeline,
      //   { $project: stageInput },
    ]);
  }
}


interface Test {
  name: string;
  email: string;
}

const agg = new BaseAggregate<Test>()
  .match({ name: "john" })
  .match({ email: "sir john" })
  .match({ email: "sir" })
  .projectActual({ email: 0, name: 1 })

I need help declaring the return type of the projectActual() function such that the return type for above example will be

interface {
name: string;
}

Thank you for your time 🙏


Solution

  • I think you want projectActual() to have the following generic call signature:

    declare projectActual<R extends Record<keyof T, 0 | 1>>(
      stageInput: R
    ): BaseAggregate<{ [K in keyof T as R[K] extends 1 ? K : never]: T[K]; }>
    

    So stageInput is of generic type R constrained to Record<keyof T, 0 | 1>, meaning that stageInput has to have every key that T has, but the values should be either 0 or 1.

    Then the type argument for the BaseAggregate output type is { [K in keyof T as R[K] extends 1 ? K : never]: T[K] }, a key-remapped mapped type that for each key K, checks if R[K] is 1. If so, then the property is output as-is, otherwise it is suppressed (if you remap a key to never, it is omitted from the resulting type). This is basically a Pick<T, Keys> where Keys is the set of keys of R which have a value type of 1.


    Let's test it out:

     interface Test {
      name: string;
      email: string;
    }
    
    const agg = new BaseAggregate<Test>()
      .match({ name: "john" })
      .match({ email: "sir john" })
      .match({ email: "sir" })
      .projectActual({ email: 0, name: 1 })
    
    /* const agg: BaseAggregate<{
        name: string;
    }> */
    

    Looks good. agg is of type BaseAggregate<{name: string}>; the email property has been suppressed effectively.

    Playground link to code