From d69cdaba975cdd2d0feaff21078864f3902f309b Mon Sep 17 00:00:00 2001
From: James Coglan <jcoglan@gmail.com>
Date: Sat, 4 Jul 2015 19:28:49 +0100
Subject: [PATCH] Allow event listeners to be passed as normal rather than
 block arguments.

This discourages the inappropriate use of 'return' inside event
listeners. Also, we return the listener, so that the caller has a
reference to it in case they do pass a block.
---
 README.md                             | 18 +++++++++---------
 lib/websocket/driver/event_emitter.rb | 14 ++++++++++----
 lib/websocket/driver/proxy.rb         |  2 +-
 spec/websocket/driver/hybi_spec.rb    |  2 +-
 4 files changed, 21 insertions(+), 15 deletions(-)

diff --git a/README.md b/README.md
index 3e5db71..e447216 100644
--- a/README.md
+++ b/README.md
@@ -146,7 +146,7 @@ module Connection
   def initialize
     @driver = WebSocket::Driver.server(self)
 
-    @driver.on(:connect) do
+    @driver.on :connect, -> (event) do
       if WebSocket::Driver.websocket?(@driver.env)
         @driver.start
       else
@@ -154,8 +154,8 @@ module Connection
       end
     end
 
-    @driver.on(:message) { |e| @driver.text(e.data) }
-    @driver.on(:close)   { |e| close_connection_after_writing }
+    @driver.on :message, -> (e) { @driver.text(e.data) }
+    @driver.on :close,   -> (e) { close_connection_after_writing }
   end
 
   def receive_data(data)
@@ -214,7 +214,7 @@ start sending incoming data to `driver.parse(data)` as normal, and call
 ```rb
 proxy = driver.proxy('http://username:password@proxy.example.com')
 
-proxy.on :connect do
+proxy.on :connect, -> (event) do
   driver.start
 end
 ```
@@ -226,7 +226,7 @@ In the event that proxy connection fails, `proxy` will emit an `:error`. You can
 inspect the proxy's response via `proxy.status` and `proxy.headers`.
 
 ```rb
-proxy.on :error do |error|
+proxy.on :error, -> (error) do
   puts error.message
   puts proxy.status
   puts proxy.headers.inspect
@@ -274,23 +274,23 @@ Note that most of these methods are commands: if they produce data that should
 be sent over the socket, they will give this to you by calling
 `socket.write(string)`.
 
-#### `driver.on('open') { |event| }`
+#### `driver.on 'open', -> (event) { }`
 
 Sets the callback block to execute when the socket becomes open.
 
-#### `driver.on('message') { |event| }`
+#### `driver.on 'message', -> (event) { }`
 
 Sets the callback block to execute when a message is received. `event` will have
 a `data` attribute containing either a string in the case of a text message or
 an array of integers in the case of a binary message.
 
-#### `driver.on('error') { |event| }`
+#### `driver.on 'error', -> (event) { }`
 
 Sets the callback to execute when a protocol error occurs due to the other peer
 sending an invalid byte sequence. `event` will have a `message` attribute
 describing the error.
 
-#### `driver.on('close') { |event| }`
+#### `driver.on 'close', -> (event) { }`
 
 Sets the callback block to execute when the socket becomes closed. The `event`
 object has `code` and `reason` attributes.
diff --git a/lib/websocket/driver/event_emitter.rb b/lib/websocket/driver/event_emitter.rb
index f6d9933..13ae268 100644
--- a/lib/websocket/driver/event_emitter.rb
+++ b/lib/websocket/driver/event_emitter.rb
@@ -6,16 +6,20 @@ module WebSocket
         @listeners = Hash.new { |h,k| h[k] = [] }
       end
 
-      def add_listener(event, &listener)
+      def add_listener(event, callable = nil, &block)
+        listener = callable || block
         @listeners[event.to_s] << listener
+        listener
       end
 
-      def on(event, &listener)
-        add_listener(event, &listener)
+      def on(event, callable = nil, &block)
+        add_listener(event, callable, &block)
       end
 
-      def remove_listener(event, &listener)
+      def remove_listener(event, callable = nil, &block)
+        listener = callable || block
         @listeners[event.to_s].delete(listener)
+        listener
       end
 
       def remove_all_listeners(event = nil)
@@ -26,6 +30,8 @@ module WebSocket
         end
       end
 
+    private
+
       def emit(event, *args)
         @listeners[event.to_s].dup.each do |listener|
           listener.call(*args)
diff --git a/lib/websocket/driver/proxy.rb b/lib/websocket/driver/proxy.rb
index dbb857f..4f567b4 100644
--- a/lib/websocket/driver/proxy.rb
+++ b/lib/websocket/driver/proxy.rb
@@ -56,7 +56,7 @@ module WebSocket
         @headers = Headers.new(@http.headers)
 
         if @status == 200
-          emit(:connect)
+          emit(:connect, ConnectEvent.new)
         else
           message = "Can't establish a connection to the server at #{@socket.url}"
           emit(:error, ProtocolError.new(message))
diff --git a/spec/websocket/driver/hybi_spec.rb b/spec/websocket/driver/hybi_spec.rb
index 487b243..6e387dc 100644
--- a/spec/websocket/driver/hybi_spec.rb
+++ b/spec/websocket/driver/hybi_spec.rb
@@ -30,7 +30,7 @@ describe WebSocket::Driver::Hybi do
 
   let :driver do
     driver = WebSocket::Driver::Hybi.new(socket, options)
-    driver.on(:open)    { |e| @open = true }
+    driver.on :open, -> e { @open = true }
     driver.on(:message) { |e| @message += e.data }
     driver.on(:error)   { |e| @error = e }
     driver.on(:close)   { |e| @close = [e.code, e.reason] }