1 files changed, 160 insertions, 79 deletions

diff --git a/src/index.ts b/src/index.ts
index 722412b..7aa6826 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -1,86 +1,167 @@
+export type EventType = string | symbol;
+
 // An event handler can take an optional event argument
 // and should not return a value
-type EventHandler = (event?: any) => void;
-type WildCardEventHandler = (type: string, event?: any) => void;
+export type Handler<T = unknown> = (event: T) => void;
+export type WildcardHandler<T = Record<string, unknown>> = (
+	type: keyof T,
+	event: T[keyof T]
+) => void;
 
 // An array of all currently registered event handlers for a type
-type EventHandlerList = Array<EventHandler>;
-type WildCardEventHandlerList = Array<WildCardEventHandler>;
+export type EventHandlerList<T = unknown> = Array<Handler<T>>;
+export type WildCardEventHandlerList<T = Record<string, unknown>> = Array<
+	WildcardHandler<T>
+>;
+
 // A map of event types and their corresponding event handlers.
-type EventHandlerMap = {
-  '*'?: WildCardEventHandlerList;
-  [type: string]: EventHandlerList;
-};
-
-/** Nitt: functional event emitter / pubsub.
- *  @name nitt
- *  @returns {Nitt}
+export type EventHandlerMap<Events extends Record<EventType, unknown>> = Map<
+	keyof Events | '*',
+	EventHandlerList<Events[keyof Events]> | WildCardEventHandlerList<Events>
+>;
+
+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;
+
+	once<Key extends keyof Events>(type: Key, handler: Handler<Events[Key]>): void;
+	once(type: '*', handler: WildcardHandler<Events>): void;
+
+	off<Key extends keyof Events>(
+		type: Key,
+		handler?: Handler<Events[Key]>
+	): void;
+	off(type: '*', handler: WildcardHandler<Events>): void;
+
+	emit<Key extends keyof Events>(type: Key, event: Events[Key]): void;
+	emit<Key extends keyof Events>(
+		type: undefined extends Events[Key] ? Key : never
+	): void;
+
+	when<Key extends keyof Events>(type: Key): Promise<void>;
+}
+
+/**
+ * Nitt: Tiny (~200b) functional event emitter / pubsub.
+ * @name nitt
+ * @returns {Nitt}
  */
-export default function nitt(all: EventHandlerMap) {
-  all = all || Object.create(null);
-
-  return {
-    /**
-     * Register an event handler for the given type.
-     *
-     * @param  {String} type	Type of event to listen for, or `"*"` for all events
-     * @param  {Function} handler Function to call in response to given event
-     */
-    on(type: string, handler: EventHandler) {
-      (all[type] || (all[type] = [])).push(handler);
-    },
-
-    /**
-     * Register an event handler that is executed just once.
-     *
-     * @param  {String} type	Type of event to listen for, or `"*"` for any event
-     * @param  {Function} handler Function to call in response to given event
-     */
-    once(type: string, handler: EventHandler) {
-      const onceHandler = (evt: any) => {
-        handler(evt);
-        this.off(type, onceHandler);
-      };
-
-      this.on(type, onceHandler);
-    },
-
-    /**
-     * Returns a promise for a single event
-     *
-     * @param  {String} type	Type of event to listen for, or `"*"` for any event
-     * @returns {Promise<any>}
-     */
-    when(type: string): Promise<any> {
-      return new Promise<any>(resolve => this.once(type, resolve));
-    },
-
-    /**
-     * Remove an event handler for the given type.
-     *
-     * @param  {String} type	Type of event to unregister `handler` from, or `"*"`
-     * @param  {Function} handler Handler function to remove
-     */
-    off(type: string, handler: EventHandler) {
-      if (all[type]) {
-        all[type].splice(all[type].indexOf(handler) >>> 0, 1);
-      }
-    },
-
-    /**
-     * Invoke all handlers for the given type.
-     * If present, `"*"` handlers are invoked after type-matched handlers.
-     *
-     * @param {String} type  The event type to invoke
-     * @param {Any} [evt]  Any value (object is recommended and powerful), passed to each handler
-     */
-    emit(type: string, evt: any) {
-      (all[type] || []).forEach(handler => {
-        handler(evt);
-      });
-      (all['*'] || []).forEach(handler => {
-        handler(type, evt);
-      });
-    },
-  };
+export default function nitt<Events extends Record<EventType, unknown>>(
+	all?: EventHandlerMap<Events>
+): Emitter<Events> {
+	type GenericEventHandler =
+		| Handler<Events[keyof Events]>
+		| WildcardHandler<Events>;
+	all = all || new Map();
+
+	return {
+		/**
+		 * A Map of event names to registered handler functions.
+		 */
+		all,
+
+		/**
+		 * Register an event handler for the given type.
+		 * @param {string|symbol} type Type of event to listen for, or `'*'` for all events
+		 * @param {Function} handler Function to call in response to given event
+		 * @memberOf nitt
+		 */
+		on<Key extends keyof Events>(type: Key, handler: GenericEventHandler) {
+			const handlers: Array<GenericEventHandler> | undefined = all!.get(type);
+			if (handlers) {
+				handlers.push(handler);
+			} else {
+				all!.set(type, [handler] as EventHandlerList<Events[keyof Events]>);
+			}
+		},
+
+		/**
+		 * Remove an event handler for the given type.
+		 * If `handler` is omitted, all handlers of the given type are removed.
+		 * @param {string|symbol} type Type of event to unregister `handler` from (`'*'` to remove a wildcard handler)
+		 * @param {Function} [handler] Handler function to remove
+		 * @memberOf nitt
+		 */
+		off<Key extends keyof Events>(type: Key, handler?: GenericEventHandler) {
+			const handlers: Array<GenericEventHandler> | undefined = all!.get(type);
+			if (handlers) {
+				if (handler) {
+					handlers.splice(handlers.indexOf(handler) >>> 0, 1);
+				} else {
+					all!.set(type, []);
+				}
+			}
+		},
+
+		/**
+		 * Invoke all handlers for the given type.
+		 * If present, `'*'` handlers are invoked after type-matched handlers.
+		 *
+		 * Note: Manually firing '*' handlers is not supported.
+		 *
+		 * @param {string|symbol} type The event type to invoke
+		 * @param {Any} [evt] Any value (object is recommended and powerful), passed to each handler
+		 * @memberOf nitt
+		 */
+		emit<Key extends keyof Events>(type: Key, evt?: Events[Key]) {
+			let handlers = all!.get(type);
+			if (handlers) {
+				(handlers as EventHandlerList<Events[keyof Events]>)
+					.slice()
+					.map((handler) => {
+						handler(evt!);
+					});
+			}
+
+			handlers = all!.get('*');
+			if (handlers) {
+				(handlers as WildCardEventHandlerList<Events>)
+					.slice()
+					.map((handler) => {
+						handler(type, evt!);
+					});
+			}
+		},
+
+		/**
+		 * Register an event handler that is executed just once for the given type.
+		 *
+		 * @param {string|symbol} type Type of event to listen for, or `'*'` for all events
+		 * @param {Function} handler Function to call in response to given event
+		 * @memberOf nitt
+		 */
+		once<Key extends keyof Events>(type: Key, handler: GenericEventHandler) {
+			if (type === '*') {
+				const onceHandler: WildcardHandler<Events> = (type: keyof Events, event: Events[keyof Events]) => {
+					(handler as WildcardHandler<Events>)(type, event);
+					this.off('*', onceHandler);
+				};
+				this.on('*', onceHandler);
+			} else {
+				const onceHandler = (evt: Events[Key]) => {
+					(handler as Handler<Events[Key]>)(evt);
+					this.off(type, onceHandler);
+				};
+				this.on(type, onceHandler);
+			}
+		},
+
+		/**
+		 * Returns a promise for a single event
+		 *
+		 * @param  {String} type	Type of event to listen for, or `"*"` for any event
+		 * @returns {Promise<any>}
+		 */
+		when<Key extends keyof Events>(type: Key): Promise<any> {
+			return new Promise((resolve: GenericEventHandler) => {
+				if (type === '*') {
+					this.once('*', (resolve as WildcardHandler<Events>));
+				} else {
+					this.once(type, (resolve as Handler<Events[Key]>));
+				}
+			});
+		}
+	};
 }