From f73aeb3e3d1a188b0940f83b4ead0690f24b30f1 Mon Sep 17 00:00:00 2001
From: Wei Wang <wei@compass.com>
Date: Thu, 26 Jul 2018 15:49:26 -0400
Subject: [PATCH] [api] Add removeContext(context) method

---
 README.md    |  6 ++++++
 index.d.ts   |  5 +++++
 index.js     | 32 ++++++++++++++++++++++++++++++++
 test/test.js | 29 +++++++++++++++++++++++++++++
 4 files changed, 72 insertions(+)

diff --git a/README.md b/README.md
index 6ce142d..961c52b 100644
--- a/README.md
+++ b/README.md
@@ -18,6 +18,8 @@ differences:
 - The `setMaxListeners`, `getMaxListeners`, `prependListener` and
   `prependOnceListener` methods are not available.
 - Support for custom context for events so there is no need to use `fn.bind`.
+- There is an additional method named `removeContext`, which removes all listeners of the given
+  context.
 - The `removeListener` method removes all matching listeners, not only the
   first.
 
@@ -61,6 +63,9 @@ or `this` value that should be set for the emitted events. This means you no
 longer have the overhead of an event that required `fn.bind` in order to get a
 custom `this` value.
 
+In addition to that, we have added a new method: `EventEmitter.removeContext`,
+which will remove all listeners of the given context.
+
 ```js
 var EE = new EventEmitter()
   , context = { foo: 'bar' };
@@ -72,6 +77,7 @@ function emitted() {
 EE.once('event-name', emitted, context);
 EE.on('another-event', emitted, context);
 EE.removeListener('another-event', emitted, context);
+EE.removeContext(context);
 ```
 
 ### Tests and benchmarks
diff --git a/index.d.ts b/index.d.ts
index 2f86f20..19cd82e 100644
--- a/index.d.ts
+++ b/index.d.ts
@@ -47,6 +47,11 @@ declare class EventEmitter<EventTypes extends string | symbol = string | symbol>
    * Remove all listeners, or those of the specified event.
    */
   removeAllListeners(event?: EventTypes): this;
+
+  /**
+   * Removes all listeners that were added with the specified context.
+   */
+  removeContext(context?: any): this;
 }
 
 declare namespace EventEmitter {
diff --git a/index.js b/index.js
index 6ea485c..fa4d778 100644
--- a/index.js
+++ b/index.js
@@ -312,6 +312,38 @@ EventEmitter.prototype.removeAllListeners = function removeAllListeners(event) {
   return this;
 };
 
+/**
+ * Remove all listeners on a specific context.
+ *
+ * @param {*} context Only remove the listeners that have this context.
+ * @returns {EventEmitter} `this`.
+ * @public
+ */
+EventEmitter.prototype.removeContext = function removeContext(context) {
+  var eventNames = this.eventNames();
+  var totalListenerCount = 0;
+
+  for (var i = 0, eventsCount = eventNames.length; i < eventsCount; i++) {
+    var evt = eventNames[i];
+    var listeners = this._events[evt];
+
+    if (listeners.fn) listeners = [listeners];
+    for (var j = 0, events = [], listenersCount = listeners.length; j < listenersCount; j++) {
+      if (listeners[j].context !== context) events.push(listeners[j]);
+    }
+
+    if (events.length) {
+      this._events[evt] = events.length === 1 ? events[0] : events;
+      totalListenerCount += events.length;
+    } else {
+      clearEvent(this, evt);
+    }
+  }
+
+  this._eventsCount = totalListenerCount;
+  return this;
+};
+
 //
 // Alias methods names because people roll like that.
 //
diff --git a/test/test.js b/test/test.js
index e5f8aaa..95ec628 100644
--- a/test/test.js
+++ b/test/test.js
@@ -568,6 +568,35 @@ describe('EventEmitter', function tests() {
     });
   });
 
+  describe('EventEmitter#removeContext', function () {
+    it('removes all listeners for the specified context', function () {
+      var e = new EventEmitter();
+      var ctx1 = {};
+      var ctx2 = {};
+      var ctx3 = {};
+
+      e.on('aaa', function () { throw new Error('oops'); }, ctx1);
+      e.on('bbb', function () { throw new Error('oops'); }, ctx1);
+      e.on('aaa', function () { throw new Error('oops'); }, ctx2);
+      e.on('bbb', function () { throw new Error('oops'); }, ctx3);
+
+      assume(e.removeContext(ctx1)).equals(e);
+      assume(e.listeners('aaa').length).equals(1);
+      assume(e.listeners('bbb').length).equals(1);
+      assume(e._eventsCount).equals(2);
+
+      assume(e.removeContext(ctx2)).equals(e);
+      assume(e.listeners('aaa').length).equals(0);
+      assume(e.listeners('bbb').length).equals(1);
+      assume(e._eventsCount).equals(1);
+
+      assume(e.removeContext(ctx3)).equals(e);
+      assume(e.listeners('aaa').length).equals(0);
+      assume(e.listeners('bbb').length).equals(0);
+      assume(e._eventsCount).equals(0);
+    });
+  });
+
   describe('EventEmitter#eventNames', function () {
     it('returns an empty array when there are no events', function () {
       var e = new EventEmitter();