types: implement Schema.discriminator

lib/schema.js

@@ -473,9 +473,39 @@ Schema.prototype.defaultOptions = function(options) {
   return options;
 };
 
+/**
+ * Inherit a Schema by applying a discriminator on an existing Schema.
+ *
+ *
+ * ####Example:
+ *
+ *     const options = { discriminatorKey: 'kind' };
+ *
+ *     const eventSchema = new mongoose.Schema({ time: Date }, options);
+ *     const Event = mongoose.model('Event', eventSchema);
+ *
+ *     // ClickedLinkEvent is a special type of Event that has
+ *     // a URL.
+ *     const ClickedLinkEvent = Event.discriminator('ClickedLink',
+ *     new mongoose.Schema({ url: String }, options));
+ *
+ *     // When you create a generic event, it can't have a URL field...
+ *     const genericEvent = new Event({ time: Date.now(), url: 'google.com' });
+ *     assert.ok(!genericEvent.url);
+ *     // But a ClickedLinkEvent can
+ *     const clickedEvent = new ClickedLinkEvent({ time: Date.now(), url: 'google.com' });
+ *     assert.ok(clickedEvent.url);
+ *
+ * @param {String} name the name of the discriminator
+ * @param {Schema} schema the Schema of the discriminated Schema
+ * @return {Schema} the Schema instance
+ * @api public
+ */
+
 Schema.prototype.discriminator = function(name, schema) {
-  this._applyDiscriminators = {};
-  this._applyDiscriminators[name] = schema;
+  this._applyDiscriminators = Object.assign({}, this._applyDiscriminators, { [name]: schema });
+
+  return this;
 };
 
 /**

test/types/schema.test.ts

@@ -1,5 +1,5 @@
 import { Schema, Document, SchemaDefinition, Model, Types } from 'mongoose';
-import { expectType, expectNotType, expectError } from 'tsd';
+import { expectType, expectError } from 'tsd';
 
 enum Genre {
   Action,
@@ -78,7 +78,7 @@ movieSchema.index({ title: 1 }, { unique: true });
 interface IProfile {
   age: number;
 }
-interface ProfileDoc extends Document, IProfile {}
+interface ProfileDoc extends Document, IProfile { }
 const ProfileSchemaDef: SchemaDefinition<IProfile> = { age: Number };
 export const ProfileSchema = new Schema<ProfileDoc, Model<ProfileDoc>, ProfileDoc>(ProfileSchemaDef);
 
@@ -87,7 +87,7 @@ interface IUser {
   profile: ProfileDoc;
 }
 
-interface UserDoc extends Document, IUser {}
+interface UserDoc extends Document, IUser { }
 
 const ProfileSchemaDef2: SchemaDefinition<IProfile> = {
   age: Schema.Types.Number
@@ -318,4 +318,11 @@ function gh10900(): void {
   const patientSchema = new Schema<IUserProp>({
     menuStatus: { type: Schema.Types.Mixed, default: {} }
   });
-}
+}
+
+// discriminator
+const eventSchema = new Schema<{ message: string }>({ message: String }, { discriminatorKey: 'kind' });
+const batchSchema = new Schema<{ name: string }>({ name: String }, { discriminatorKey: 'kind' });
+const discriminatedSchema = batchSchema.discriminator('event', eventSchema);
+
+expectType<Schema<Omit<{ name: string }, 'message'> & { message: string }>>(discriminatedSchema);

types/index.d.ts

@@ -120,6 +120,17 @@ declare module 'mongoose' {
     useProjection?: boolean;
   }
 
+  export type DiscriminatorModel<M, T> = T extends Model<infer T1, infer T2, infer T3, infer T4>
+    ?
+    M extends Model<infer M1, infer M2, infer M3, infer M4>
+      ? Model<Omit<M1, keyof T1> & T1, M2 | T2, M3 | T3, M4 | T4>
+      : M
+    : M;
+
+  export type DiscriminatorSchema<DocType, M, TInstanceMethods, TQueryHelpers, TVirtuals, T> = T extends Schema<infer T1, infer T2, infer T3, infer T4, infer T5>
+    ? Schema<Omit<DocType, keyof T1> & T1, DiscriminatorModel<T2, M>, T3 | TInstanceMethods, T4 | TQueryHelpers, T5 | TVirtuals>
+    : Schema<DocType, M, TInstanceMethods, TQueryHelpers, TVirtuals>;
+
   export class Schema<DocType = any, M = Model<DocType, any, any, any>, TInstanceMethods = {}, TQueryHelpers = {}, TVirtuals = any> extends events.EventEmitter {
     /**
      * Create a new schema
@@ -142,6 +153,8 @@ declare module 'mongoose' {
     /** Returns a copy of this schema */
     clone<T = this>(): T;
 
+    discriminator<T extends Schema = Schema>(name: string, schema: T): DiscriminatorSchema<DocType, M, TInstanceMethods, TQueryHelpers, TVirtuals, T>;
+
     /** Returns a new schema that has the picked `paths` from this schema. */
     pick<T = this>(paths: string[], options?: SchemaOptions): T;
 
@@ -474,10 +487,10 @@ declare module 'mongoose' {
   type LeanTypeOrArray<T> = T extends unknown[] ? LeanArray<T> : LeanType<T>;
 
   export type LeanArray<T extends unknown[]> =
-  // By checking if it extends Types.Array we can get the original base type before collapsing down,
-  // rather than trying to manually remove the old types. This matches both Array and DocumentArray
+    // By checking if it extends Types.Array we can get the original base type before collapsing down,
+    // rather than trying to manually remove the old types. This matches both Array and DocumentArray
     T extends Types.Array<infer U> ? LeanTypeOrArray<U>[] :
-      // If it isn't a custom mongoose type we fall back to "do our best"
+    // If it isn't a custom mongoose type we fall back to "do our best"
       T extends unknown[][] ? LeanArray<T[number]>[] : LeanType<T[number]>[];
 
   export type _LeanDocument<T> = {