svn commit: r920660 - /incubator/thrift/trunk/lib/py/src/transport/THttpClient.py

[dr...@apache.org]

Author: dreiss
Date: Tue Mar  9 05:19:16 2010
New Revision: 920660

URL: http://svn.apache.org/viewvc?rev=920660&view=rev
Log:
THRIFT-658. python: Add THttpClient.setTimeout()

Modified:
    incubator/thrift/trunk/lib/py/src/transport/THttpClient.py

Modified: incubator/thrift/trunk/lib/py/src/transport/THttpClient.py
URL: http://svn.apache.org/viewvc/incubator/thrift/trunk/lib/py/src/transport/THttpClient.py?rev=920660&r1=920659&r2=920660&view=diff
==============================================================================
--- incubator/thrift/trunk/lib/py/src/transport/THttpClient.py (original)
+++ incubator/thrift/trunk/lib/py/src/transport/THttpClient.py Tue Mar  9 05:19:16 2010
@@ -23,6 +23,7 @@ from cStringIO import StringIO
 import urlparse
 import httplib
 import warnings
+import socket
 
 class THttpClient(TTransportBase):
 
@@ -55,6 +56,7 @@ class THttpClient(TTransportBase):
       self.path = parsed.path
     self.__wbuf = StringIO()
     self.__http = None
+    self.__timeout = None
 
   def open(self):
     if self.scheme == 'http':
@@ -69,12 +71,30 @@ class THttpClient(TTransportBase):
   def isOpen(self):
     return self.__http != None
 
+  def setTimeout(self, ms):
+    if not hasattr(socket, 'getdefaulttimeout'):
+      raise NotImplementedError
+
+    if ms is None:
+      self.__timeout = None
+    else:
+      self.__timeout = ms/1000.0
+
   def read(self, sz):
     return self.__http.file.read(sz)
 
   def write(self, buf):
     self.__wbuf.write(buf)
 
+  def __withTimeout(f):
+    def _f(*args, **kwargs):
+      orig_timeout = socket.getdefaulttimeout()
+      socket.setdefaulttimeout(args[0].__timeout)
+      result = f(*args, **kwargs)
+      socket.setdefaulttimeout(orig_timeout)
+      return result
+    return _f
+
   def flush(self):
     if self.isOpen():
       self.close()
@@ -98,3 +118,7 @@ class THttpClient(TTransportBase):
 
     # Get reply to flush the request
     self.code, self.message, self.headers = self.__http.getreply()
+
+  # Decorate if we know how to timeout
+  if hasattr(socket, 'getdefaulttimeout'):
+    flush = __withTimeout(flush)