Merge pull request #4668 from rockmacaca/class-models

Integration for ES6 classes

lib/index.js

@@ -309,14 +309,23 @@ Mongoose.prototype.disconnect.$hasSideEffects = true;
  *     var collectionName = 'actor'
  *     var M = mongoose.model('Actor', schema, collectionName)
  *
- * @param {String} name model name
+ * @param {String|Function} name model name or class extending Model
  * @param {Schema} [schema]
- * @param {String} [collection] name (optional, induced from model name)
+ * @param {String} [collection] name (optional, inferred from model name)
  * @param {Boolean} [skipInit] whether to skip initialization (defaults to false)
  * @api public
  */
 
 Mongoose.prototype.model = function(name, schema, collection, skipInit) {
+  var model;
+  if (typeof name === 'function') {
+    model = name;
+    name = model.name;
+    if (!(model.prototype instanceof Model)) {
+      throw new mongoose.Error('The provided class ' + name + ' must extend Model');
+    }
+  }
+
   if (typeof schema === 'string') {
     collection = schema;
     schema = false;
@@ -358,7 +367,6 @@ Mongoose.prototype.model = function(name, schema, collection, skipInit) {
     this._applyPlugins(schema);
   }
 
-  var model;
   var sub;
 
   // connection.model() may be passing a different schema for
@@ -397,7 +405,7 @@ Mongoose.prototype.model = function(name, schema, collection, skipInit) {
   }
 
   var connection = options.connection || this.connection;
-  model = this.Model.compile(name, schema, collection, connection, this);
+  model = this.Model.compile(model || name, schema, collection, connection, this);
 
   if (!skipInit) {
     model.init();

lib/model.js

@@ -3198,7 +3198,7 @@ Model._getSchema = function _getSchema(path) {
 /*!
  * Compiler utility.
  *
- * @param {String} name model name
+ * @param {String|Function} name model name or class extending Model
  * @param {Schema} schema
  * @param {String} collectionName
  * @param {Connection} connection
@@ -3215,19 +3215,28 @@ Model.compile = function compile(name, schema, collectionName, connection, base)
     schema.add(o);
   }
 
-  // generate new class
-  function model(doc, fields, skipId) {
-    if (!(this instanceof model)) {
-      return new model(doc, fields, skipId);
-    }
-    Model.call(this, doc, fields, skipId);
+  var model;
+  if (typeof name === 'function' && name.prototype instanceof Model) {
+    model = name;
+    name = model.name;
+    schema.loadClass(model, true);
+  } else {
+    // generate new class
+    model = function model(doc, fields, skipId) {
+      if (!(this instanceof model)) {
+        return new model(doc, fields, skipId);
+      }
+      Model.call(this, doc, fields, skipId);
+    };
   }
 
   model.hooks = schema.s.hooks.clone();
   model.base = base;
   model.modelName = name;
-  model.__proto__ = Model;
-  model.prototype.__proto__ = Model.prototype;
+  if (!(model.prototype instanceof Model)) {
+    model.__proto__ = Model;
+    model.prototype.__proto__ = Model.prototype;
+  }
   model.model = Model.prototype.model;
   model.db = model.prototype.db = connection;
   model.discriminators = model.prototype.discriminators = undefined;

lib/schema.js

@@ -1543,6 +1543,49 @@ Schema.prototype.remove = function(path) {
   }
 };
 
+/**
+ * Loads an ES6 class into a schema. Maps setters + getters, static methods, and instance methods to schema virtuals, statics, and methods.
+ *
+ * @param {Function} model
+ */
+Schema.prototype.loadClass = function(model, virtualsOnly) {
+  if (model === Object.prototype || model === Function.prototype) {
+    return this;
+  }
+
+  // Add static methods
+  if (!virtualsOnly) {
+    Object.getOwnPropertyNames(model).forEach(function(name) {
+      if (name.match(/^(length|name|prototype)$/)) {
+        return;
+      }
+      var method = Object.getOwnPropertyDescriptor(model, name);
+      if (typeof method.value === 'function') this.static(name, method.value);
+    }, this);
+  }
+
+  // Add methods and virtuals
+  Object.getOwnPropertyNames(model.prototype).forEach(function(name) {
+    if (name.match(/^(constructor)$/)) {
+      return;
+    }
+    var method = Object.getOwnPropertyDescriptor(model.prototype, name);
+    if (!virtualsOnly) {
+      if (typeof method.value === 'function') {
+        this.method(name, method.value);
+      }
+    }
+    if (typeof method.get === 'function') {
+      this.virtual(name).get(method.get);
+    }
+    if (typeof method.set === 'function') {
+      this.virtual(name).set(method.set);
+    }
+  }, this);
+
+  return (this.loadClass(Object.getPrototypeOf(model)));
+};
+
 /*!
  * ignore
  */