Merge pull request #2480 from alabid/gh-2167

gh-2167 give discriminators access to their own instance or class methods

Author: vkarpov15

lib/model.js

@@ -710,6 +710,10 @@ Model.discriminator = function discriminator (name, schema) {
   this.discriminators[name] = this.db.model(name, schema, this.collection.name);
   this.discriminators[name].prototype.__proto__ = this.prototype;
 
+  // apply methods and statics
+  applyMethods(this.discriminators[name], schema);
+  applyStatics(this.discriminators[name], schema);
+
   return this.discriminators[name];
 };
 
@@ -2545,13 +2549,9 @@ Model.compile = function compile (name, schema, collectionName, connection, base
     , collectionOptions
   );
 
-  // apply methods
-  for (var i in schema.methods)
-    model.prototype[i] = schema.methods[i];
-
-  // apply statics
-  for (var i in schema.statics)
-    model[i] = schema.statics[i];
+  // apply methods and statics
+  applyMethods(model, schema);
+  applyStatics(model, schema);
 
   model.schema = model.prototype.schema;
   model.options = model.prototype.options;
@@ -2560,6 +2560,53 @@ Model.compile = function compile (name, schema, collectionName, connection, base
   return model;
 };
 
+/*!
+ * Register methods for this model
+ *
+ * @param {Model} model
+ * @param {Schema} schema
+ */
+var applyMethods = function(model, schema) {
+  for (var i in schema.methods) {
+    if (typeof schema.methods[i] === 'function') {
+      model.prototype[i] = schema.methods[i];
+    } else {
+      (function(_i) {
+        Object.defineProperty(model.prototype, _i, {
+          get: function() {
+            var h = {};
+            for (var k in schema.methods[_i]) {
+              h[k] = schema.methods[_i][k].bind(this);
+            }
+            return h;
+          }
+        });
+      })(i);
+    }
+  }
+
+  for (var i in schema.methods)
+    model.prototype[i] = schema.methods[i];
+};
+
+/*!
+ * Register statics for this model
+ * @param {Model} model
+ * @param {Schema} schema
+ */
+var applyStatics = function(model, schema) {
+  for (var i in schema.statics) {
+    // use defineProperty so that static props can't be overwritten
+    Object.defineProperty(model, i, {
+      value: schema.statics[i],
+      writable: false
+    });
+  }
+
+  for (var i in schema.statics)
+    model[i] = schema.statics[i];
+};
+
 /*!
  * Subclass this model with `conn`, `schema`, and `collection` settings.
  *

test/model.discriminator.test.js

@@ -6,6 +6,7 @@ var start = require('./common')
   , mongoose = start.mongoose
   , Schema = mongoose.Schema
   , assert = require('assert')
+  , util = require('util')
   , clone = require('../lib/utils').clone
   , random = require('../lib/utils').random;
 
@@ -83,6 +84,36 @@ describe('model', function() {
       done();
     });
 
+    it('can define static and instance methods', function(done) {
+      function BossBaseSchema() {
+        Schema.apply(this, arguments);
+
+        this.add({
+          name: String,
+          createdAt: Date
+        });
+      }
+      util.inherits(BossBaseSchema, Schema);
+
+      var PersonSchema = new BossBaseSchema();
+      var BossSchema = new BossBaseSchema({ department: String });
+      BossSchema.methods.myName = function(){
+        return this.name;
+      };
+      BossSchema.statics.currentPresident = function(){
+        return 'obama';
+      };
+      var Person = db.model('Person', PersonSchema);
+      var Boss = Person.discriminator('Boss', BossSchema);
+
+      var boss = new Boss({name:'Bernenke'});
+      assert.equal(boss.myName(), 'Bernenke');
+      assert.equal(boss.notInstanceMethod, undefined);
+      assert.equal(Boss.currentPresident(), 'obama');
+      assert.equal(Boss.notStaticMethod, undefined);
+      done();
+    });
+
     it('sets schema root discriminator mapping', function(done) {
       assert.deepEqual(PersonSchema.discriminatorMapping, { key: '__t', value: null, isRoot: true });
       done();