From 502c80be05baa59ea126ad2ba53166e1723e05bb Mon Sep 17 00:00:00 2001
From: liris <liris.pp@gmail.com>
Date: Mon, 4 Mar 2013 09:37:42 +0900
Subject: [PATCH] - fixed #29 - Passing socketopt parameter to Websocket ctor.

---
 setup.py          |  2 +-
 test_websocket.py | 12 ++++++++++++
 websocket.py      | 18 +++++++++++++-----
 3 files changed, 26 insertions(+), 6 deletions(-)

diff --git a/setup.py b/setup.py
index fdd978d..2e5c030 100644
--- a/setup.py
+++ b/setup.py
@@ -1,6 +1,6 @@
 from setuptools import setup
 
-VERSION = "0.9.0"
+VERSION = "0.10.0"
 
 
 setup(
diff --git a/test_websocket.py b/test_websocket.py
index b2726a4..a706a0d 100644
--- a/test_websocket.py
+++ b/test_websocket.py
@@ -4,6 +4,7 @@
 import base64
 import uuid
 import unittest
+import socket
 
 # websocket-client
 import websocket as ws
@@ -32,6 +33,7 @@ class StringSockMock:
 
     def send(self, data):
         self.sent.append(data)
+        return len(data)
 
 
 class HeaderSockMock(StringSockMock):
@@ -304,5 +306,15 @@ class WebSocketAppTest(unittest.TestCase):
         # Note: We can't use 'is' for comparing the functions directly, need to use 'id'.
         self.assertEquals(WebSocketAppTest.get_mask_key_id, id(my_mask_key_func))
 
+
+class SockOptTest(unittest.TestCase):
+    def testSockOpt(self):
+        sockopt = ((socket.IPPROTO_TCP, socket.TCP_NODELAY, 1),)
+        s = ws.WebSocket(sockopt = sockopt)
+        self.assertNotEquals(s.sock.getsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY), 0)
+        s = ws.create_connection("ws://echo.websocket.org", sockopt = sockopt)
+        self.assertNotEquals(s.sock.getsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY), 0)
+
+
 if __name__ == "__main__":
     unittest.main()
diff --git a/websocket.py b/websocket.py
index ad733a1..3063554 100644
--- a/websocket.py
+++ b/websocket.py
@@ -172,7 +172,8 @@ def create_connection(url, timeout=None, **options):
     options: current support option is only "header".
              if you set header as dict value, the custom HTTP headers are added.
     """
-    websock = WebSocket()
+    sockopt = options.get("sockopt", ())
+    websock = WebSocket(sockopt=sockopt)
     websock.settimeout(timeout != None and timeout or default_timeout)
     websock.connect(url, **options)
     return websock
@@ -349,14 +350,18 @@ class WebSocket(object):
 
     get_mask_key: a callable to produce new mask keys, see the set_mask_key
       function's docstring for more details
+    sockopt: values for socket.setsockopt.
+        sockopt must be tuple and each element is argument of sock.setscokopt.
     """
 
-    def __init__(self, get_mask_key = None):
+    def __init__(self, get_mask_key = None, sockopt = ()):
         """
         Initalize WebSocket object.
         """
         self.connected = False
         self.io_sock = self.sock = socket.socket()
+        for opts in sockopt:
+            self.sock.setsockopt(*opts)
         self.get_mask_key = get_mask_key
 
     def set_mask_key(self, func):
@@ -687,7 +692,8 @@ class WebSocketApp(object):
     """
     def __init__(self, url, header = [],
                  on_open = None, on_message = None, on_error = None,
-                 on_close = None, keep_running = True, get_mask_key = None):
+                 on_close = None, keep_running = True, get_mask_key = None,
+                 sockopt=()):
         """
         url: websocket url.
         header: custom header for websocket handshake.
@@ -734,15 +740,17 @@ class WebSocketApp(object):
         self.keep_running = False
         self.sock.close()
 
-    def run_forever(self):
+    def run_forever(self, sockopt=()):
         """
         run event loop for WebSocket framework.
         This loop is infinite loop and is alive during websocket is available.
+        sockopt: values for socket.setsockopt.
+            sockopt must be tuple and each element is argument of sock.setscokopt.
         """
         if self.sock:
             raise WebSocketException("socket is already opened")
         try:
-            self.sock = WebSocket(self.get_mask_key)
+            self.sock = WebSocket(self.get_mask_key, sockopt = sockopt)
             self.sock.connect(self.url, header = self.header)
             self._run_with_no_err(self.on_open)
             while self.keep_running: