From 08f357d5694800e7359ab64fbcdf16366dd70922 Mon Sep 17 00:00:00 2001
From: Adam Smith <black8192water@gmail.com>
Date: Sat, 14 Dec 2024 12:54:27 +0300
Subject: [PATCH] feat: added a signal for unsubscribing on abort

---
 src/index.ts       | 30 +++++++++++++++++++++++++++---
 test/index_test.ts | 37 +++++++++++++++++++++++++++++++++++++
 2 files changed, 64 insertions(+), 3 deletions(-)

diff --git a/src/index.ts b/src/index.ts
index 17672aa..ef08b73 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -20,11 +20,23 @@ export type EventHandlerMap<Events extends Record<EventType, unknown>> = Map<
 	EventHandlerList<Events[keyof Events]> | WildCardEventHandlerList<Events>
 >;
 
+export type SubscribeParams = {
+	signal: AbortSignal
+}
+
 export interface Emitter<Events extends Record<EventType, unknown>> {
 	all: EventHandlerMap<Events>;
 
-	on<Key extends keyof Events>(type: Key, handler: Handler<Events[Key]>): void;
-	on(type: '*', handler: WildcardHandler<Events>): void;
+	on<Key extends keyof Events>(
+		type: Key,
+		handler: Handler<Events[Key]>,
+		params?: Partial<SubscribeParams>
+	): void;
+	on(
+		type: '*',
+		handler: WildcardHandler<Events>,
+		params?: Partial<SubscribeParams>
+	): void;
 
 	off<Key extends keyof Events>(
 		type: Key,
@@ -63,13 +75,25 @@ export default function mitt<Events extends Record<EventType, unknown>>(
 		 * @param {Function} handler Function to call in response to given event
 		 * @memberOf mitt
 		 */
-		on<Key extends keyof Events>(type: Key, handler: GenericEventHandler) {
+		on<Key extends keyof Events>(
+			type: Key,
+			handler: GenericEventHandler,
+			params?: Partial<SubscribeParams>,
+		) {
+			if (params?.signal?.aborted) {
+				return;
+			}
+			
 			const handlers: Array<GenericEventHandler> | undefined = all!.get(type);
 			if (handlers) {
 				handlers.push(handler);
 			} else {
 				all!.set(type, [handler] as EventHandlerList<Events[keyof Events]>);
 			}
+
+			params?.signal?.addEventListener('abort', () => {
+				this.off<Key>(type, handler as Handler<Events[keyof Events]>);
+			});
 		},
 
 		/**
diff --git a/test/index_test.ts b/test/index_test.ts
index 31c9495..4ea4a52 100644
--- a/test/index_test.ts
+++ b/test/index_test.ts
@@ -103,6 +103,43 @@ describe('mitt#', () => {
 			inst.on('foo', foo);
 			expect(events.get('foo')).to.deep.equal([foo, foo]);
 		});
+
+		it('should unsubscribe when signal aborted', () => {
+			const foo = spy();
+			const abortController = new AbortController();
+
+			inst.on(
+				'foo',
+				foo,
+				{ signal: abortController.signal }
+			);
+			inst.emit('foo');
+
+			abortController.abort();
+
+			inst.emit('foo');
+
+			expect(foo).to.have.been.callCount(1);
+		});
+
+		it('should make multiple unsubscribe when aborted', () => {
+			const foo = spy();
+			const bar = spy();
+			const abortController = new AbortController();
+
+			inst.on('foo', foo, { signal: abortController.signal });
+			inst.on('bar', bar, { signal: abortController.signal });
+			inst.emit('foo');
+			inst.emit('bar');
+
+			abortController.abort();
+
+			inst.emit('foo');
+			inst.emit('bar');
+
+			expect(foo).to.have.been.callCount(1);
+			expect(bar).to.have.been.callCount(1);
+		});
 	});
 
 	describe('off()', () => {