Patch WebSocketApp class to make it inheritable

4 files changed, 127 insertions, 12 deletions

diff --git a/README.rst b/README.rst
index 86e3235..aca0da0 100644
--- a/README.rst
+++ b/README.rst
@@ -106,6 +106,54 @@ This example is similar to how WebSocket code looks in browsers using JavaScript
         ws.run_forever()
 
 
+Class inheritance
+--------------------------------
+Same as above with class inheritance.
+
+.. code:: python
+
+import websocket
+    from threading import Thread
+    import time
+    import sys
+
+
+    class MyApp(websocket.WebSocketApp):
+        def on_message(self, message):
+            print(message)
+
+        def on_error(self, error):
+            print(error)
+
+        def on_close(self):
+            print("### closed ###")
+
+        def on_open(self):
+            def run(*args):
+                for i in range(3):
+                    # send the message, then wait
+                    # so thread doesn't exit and socket
+                    # isn't closed
+                    self.send("Hello %d" % i)
+                    time.sleep(1)
+
+                time.sleep(1)
+                self.close()
+                print("Thread terminating...")
+
+            Thread(target=run).start()
+
+
+    if __name__ == "__main__":
+        websocket.enableTrace(True)
+        if len(sys.argv) < 2:
+            host = "ws://echo.websocket.org/"
+        else:
+            host = sys.argv[1]
+        ws = MyApp(host)
+        ws.run_forever()
+
+
 Short-lived one-off send-receive
 --------------------------------
 This is if you want to communicate a short message and disconnect immediately when done.
diff --git a/examples/echoapp_client_inheritance.py b/examples/echoapp_client_inheritance.py
new file mode 100644
index 0000000..775e632
--- /dev/null
+++ b/examples/echoapp_client_inheritance.py
@@ -0,0 +1,40 @@
+import websocket
+from threading import Thread
+import time
+import sys
+
+
+class MyApp(websocket.WebSocketApp):
+    def on_message(self, message):
+        print(message)
+
+    def on_error(self, error):
+        print(error)
+
+    def on_close(self):
+        print("### closed ###")
+
+    def on_open(self):
+        def run(*args):
+            for i in range(3):
+                # send the message, then wait
+                # so thread doesn't exit and socket
+                # isn't closed
+                self.send("Hello %d" % i)
+                time.sleep(1)
+
+            time.sleep(1)
+            self.close()
+            print("Thread terminating...")
+
+        Thread(target=run).start()
+
+
+if __name__ == "__main__":
+    websocket.enableTrace(True)
+    if len(sys.argv) < 2:
+        host = "ws://echo.websocket.org/"
+    else:
+        host = sys.argv[1]
+    ws = MyApp(host)
+    ws.run_forever()
diff --git a/websocket/_app.py b/websocket/_app.py
index ecdc916..f31ad8f 100644
--- a/websocket/_app.py
+++ b/websocket/_app.py
@@ -23,6 +23,7 @@ Copyright (C) 2010 Hiroki Ohtani(liris)
 """
 WebSocketApp provides higher level APIs.
 """
+import inspect
 import select
 import sys
 import threading
@@ -125,16 +126,17 @@ class WebSocketApp(object):
         self.url = url
         self.header = header if header is not None else []
         self.cookie = cookie
-        self.on_open = on_open
-        self.on_message = on_message
-        self.on_data = on_data
-        self.on_error = on_error
-        self.on_close = on_close
-        self.on_ping = on_ping
-        self.on_pong = on_pong
-        self.on_cont_message = on_cont_message
+
+        self.on_open = on_open or getattr(self, 'on_open', None)
+        self.on_message = on_message or getattr(self, 'on_message', None)
+        self.on_data = on_data or getattr(self, 'on_data', None)
+        self.on_error = on_error or getattr(self, 'on_error', None)
+        self.on_close = on_close or getattr(self, 'on_close', None)
+        self.on_ping = on_ping or getattr(self, 'on_ping', None)
+        self.on_pong = on_pong or getattr(self, 'on_pong', None)
+        self.on_cont_message = on_cont_message or getattr(self, 'on_cont_message', None)
+        self.get_mask_key = get_mask_key or getattr(self, 'get_mask_key', None)
         self.keep_running = False
-        self.get_mask_key = get_mask_key
         self.sock = None
         self.last_ping_tm = 0
         self.last_pong_tm = 0
@@ -317,7 +319,10 @@ class WebSocketApp(object):
     def _callback(self, callback, *args):
         if callback:
             try:
-                callback(self, *args)
+                if inspect.ismethod(callback):
+                    callback(*args)
+                else:
+                    callback(self, *args)
             except Exception as e:
                 _logging.error("error from callback {}: {}".format(callback, e))
                 if _logging.isEnabledForDebug():
diff --git a/websocket/tests/test_websocket.py b/websocket/tests/test_websocket.py
index 90399f5..a090cd6 100644
--- a/websocket/tests/test_websocket.py
+++ b/websocket/tests/test_websocket.py
@@ -542,6 +542,28 @@ class WebSocketAppTest(unittest.TestCase):
         # Note: We can't use 'is' for comparing the functions directly, need to use 'id'.
         # self.assertEqual(WebSocketAppTest.get_mask_key_id, id(my_mask_key_func))
 
+    def testSettingClassCallbacks(self):
+        """ App class should provide possibility to set callback functions via class instantiate call, class inheritance
+        and via setting instance attributes
+        """
+        class TestApp(ws.WebSocketApp):
+            def on_close(self):
+                pass
+
+        def on_open(ws):
+            pass
+
+        def on_message(ws):
+            pass
+
+        app = TestApp('ws://www.example.com/', on_open=on_open)
+        app.on_message = on_message
+
+        #assert isinstance(app.on_close, function)
+        assert callable(app.on_open)
+        assert callable(app.on_message)
+        assert callable(app.on_close)
+
 
 class SockOptTest(unittest.TestCase):
     @unittest.skipUnless(TEST_WITH_INTERNET, "Internet-requiring tests are disabled")
@@ -551,6 +573,7 @@ class SockOptTest(unittest.TestCase):
         self.assertNotEqual(s.sock.getsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY), 0)
         s.close()
 
+
 class UtilsTest(unittest.TestCase):
     def testUtf8Validator(self):
         state = validate_utf8(six.b('\xf0\x90\x80\x80'))
@@ -560,6 +583,7 @@ class UtilsTest(unittest.TestCase):
         state = validate_utf8(six.b(''))
         self.assertEqual(state, True)
 
+
 class ProxyInfoTest(unittest.TestCase):
     def setUp(self):
         self.http_proxy = os.environ.get("http_proxy", None)
@@ -580,7 +604,6 @@ class ProxyInfoTest(unittest.TestCase):
         elif "https_proxy" in os.environ:
             del os.environ["https_proxy"]
 
-
     def testProxyFromArgs(self):
         self.assertEqual(get_proxy_info("echo.websocket.org", False, proxy_host="localhost"), ("localhost", 0, None))
         self.assertEqual(get_proxy_info("echo.websocket.org", False, proxy_host="localhost", proxy_port=3128), ("localhost", 3128, None))
@@ -601,7 +624,6 @@ class ProxyInfoTest(unittest.TestCase):
         self.assertEqual(get_proxy_info("echo.websocket.org", True, proxy_host="localhost", proxy_port=3128, no_proxy=["echo.websocket.org"], proxy_auth=("a", "b")),
             (None, 0, None))
 
-
     def testProxyFromEnv(self):
         os.environ["http_proxy"] = "http://localhost/"
         self.assertEqual(get_proxy_info("echo.websocket.org", False), ("localhost", None, None))