@@ -12,8 +12,8 @@ import urlparse
import urllib
import pycurl
import cStringIO
import Queue
import threading
import nilmdb.iteratorizer
class Error(Exception):
"""Base exception for both ClientError and ServerError responses"""
@@ -89,29 +89,49 @@ class HTTPClient(object):
else:
raise Error(**args)
def _perform_generator(self):
"""Like self.curl.perform(), but runs the perform in a
separate thread and returns a generator instead of using
the WRITEFUNCTION callback"""
queue = Queue.Queue(maxsize=1)
def callback(data):
queue.put((1, data))
def thread():
try:
self.curl.perform()
except:
info = sys.exc_info()
queue.put((2, info))
queue.put((0, None))
self.curl.setopt(pycurl.WRITEFUNCTION, callback)
self.curl.perform()
def _req_generator(self, url, params):
"""
Like self._req(), but runs the perform in a separate thread.
It returns a generator that spits out arbitrary-sized chunks
of the resulting data, instead of using the WRITEFUNCTION
callback.
"""
self._setup_url(url, params)
self._status = None
error_body = ""
def header_callback(data):
if self._status is None:
self._status = int(data.split(" ")[1])
self.curl.setopt(pycurl.HEADERFUNCTION, header_callback)
def func(callback):
self.curl.setopt(pycurl.WRITEFUNCTION, callback)
self.curl.perform()
try:
for i in nilmdb.iteratorizer.Iteratorizer(func):
if self._status == 200:
# If we had a 200 response, yield the data to the caller.
yield i
else:
# Otherwise, collect it into an error string.
error_body += i
except pycurl.error as e:
raise ServerError(status = "502 Error",
url = self.url,
message = e[1])
# Raise an exception if there was an error
self._check_error(error_body)
def _req(self, url, params):
"""GET or POST that returns raw data"""
"""
GET or POST that returns raw data. Returns the body
data as a string, or raises an error if it contained an error.
"""
self._setup_url(url, params)
body = cStringIO.StringIO()
self.curl.setopt(pycurl.WRITEFUNCTION, body.write)
def header_callback(data):
pass
self.curl.setopt(pycurl.HEADERFUNCTION, header_callback)
try:
self.curl.perform()
except pycurl.error as e:
@@ -119,27 +139,79 @@ class HTTPClient(object):
url = self.url,
message = e[1])
body_str = body.getvalue()
# Raise an exception if there was an error
self._check_error(body_str)
return body_str
def close(self):
self.curl.close()
def get(self, url, params = None, retjson = True):
"""Simple GET"""
self.curl.setopt(pycurl.UPLOAD, 0)
def _iterate_lines(self, it):
"""
Given an iterator that returns arbitrarily-sized chunks
of data, return '\n'-delimited lines of text
"""
partial = ""
for chunk in it:
partial += chunk
lines = partial.split("\n")
for line in lines[0:-1]:
yield line
partial = lines[-1]
if partial != "":
yield partial
# Non-generator versions
def _doreq(self, url, params, retjson):
"""
Perform a request, and return the body.
url: URL to request (relative to baseurl)
params: dictionary of query parameters
retjson: expect JSON and return python objects instead of string
"""
out = self._req(url, params)
if retjson:
return json.loads(out)
return out
def get(self, url, params = None, retjson = True):
"""Simple GET"""
self.curl.setopt(pycurl.UPLOAD, 0)
return self._doreq(url, params, retjson)
def put(self, url, postdata, params = None, retjson = True):
"""Simple PUT"""
self._setup_url(url, params)
data = cStringIO.StringIO(postdata)
self.curl.setopt(pycurl.UPLOAD, 1)
self.curl.setopt(pycurl.READFUNCTION, data.read)
out = self._req(url, params)
if retjson:
return json.loads(out)
return out
return self._doreq(url, params, retjson)
# Generator versions
def _doreq_gen(self, url, params, retjson):
"""
Perform a request, and return lines of the body in a generator.
url: URL to request (relative to baseurl)
params: dictionary of query parameters
retjson: expect JSON and yield python objects instead of strings
"""
for line in self._iterate_lines(self._req_generator(url, params)):
if retjson:
yield json.loads(line)
else:
yield line
def get_gen(self, url, params = None, retjson = True):
"""Simple GET, returning a generator"""
self.curl.setopt(pycurl.UPLOAD, 0)
return self._doreq_gen(url, params, retjson)
def put_gen(self, url, postdata, params = None, retjson = True):
"""Simple PUT, returning a generator"""
self._setup_url(url, params)
data = cStringIO.StringIO(postdata)
self.curl.setopt(pycurl.UPLOAD, 1)
self.curl.setopt(pycurl.READFUNCTION, data.read)
return self._doreq_gen(url, params, retjson)