Python:Mwclient/http.py

import urllib2
import cookielib
import urlparse
import httplib
import socket
 
import errors
 
__ver__ = '0.4.0'
 
class Request(urllib2.Request):
	def __init__(self, url, data=None, headers={},
		origin_req_host=None, unverifiable=False):
		urllib2.Request.__init__(self, url, data, headers, origin_req_host, unverifiable)
 
		self.add_header('User-Agent', 'MwClient-' + __ver__)
 
class HTTPConnection(object):
	def __init__(self, host):
		self.host = host
		self.cookies = cookielib.CookieJar()
 
	def __request(self, method, path, headers, data):
		#ignore method
		path = 'http://%s%s' % (self.host, path)	
 
		req = Request(path, data, headers)
		self.cookies.add_cookie_header(req)
		res = urllib2.urlopen(req)
		self.cookies.extract_cookies(res, req)
		return res
 
	def get(self, path, headers = {}, data = None):
		return self.__request('GET', path, headers, data)
	post = get
 
class CookieJar(dict):
	def __init__(self):
		dict.__init__(self, ())
	def extract_cookies(self, response):
		for cookie in response.msg.getallmatchingheaders('Set-Cookie'):
			self.parse_cookie(cookie.strip())
		if response.getheader('set-cookie2', None):
			# ...
			raise RuntimeError, 'Set-Cookie2', value
	def parse_cookie(self, cookie):
		if not cookie: return
		value, attrs = cookie.split(': ', 1)[1].split(';', 1)
		i = value.strip().split('=')
		if len(i) == 1 and i[0] in self: 
			del self[i[0]]
		else:
			self[i[0]] = i[1]
 
	def get_cookie_header(self):
		return '; '.join(('%s=%s' % i for i in self.iteritems()))
	def __iter__(self):
		for k, v in self.iteritems():
			yield Cookie(k, v)
 
class Cookie(object):
	def __init__(self, name, value):
		self.name = name
		self.value = value
 
class HTTPPersistentConnection(object):
	def __init__(self, host):
		self.host = host
		self.cookies = CookieJar()
		self._conn = httplib.HTTPConnection(host)
		#self._conn.set_debuglevel(100)
		self._conn.connect()
 
	def __request(self, method, path, headers, data):		
		if not headers: headers = {}
		if not data: data = ''
		headers['Connection'] = 'Keep-Alive'
		headers['User-Agent'] = 'MwClient-' + __ver__
		if self.cookies: headers['Cookie'] = self.cookies.get_cookie_header()
 
		self._conn.request(method, path, data, headers)
 
		try:
			res = self._conn.getresponse()
		except httplib.BadStatusLine:
			self._conn.close()
			self._conn.connect()
			self._conn.request(method, path, data, headers)
			res = self._conn.getresponse()
 
		if res.status >= 500:
			self._conn.request(method, path, data, headers)
			res = self._conn.getresponse()
 
		self.cookies.extract_cookies(res)
		#print 'Cookies:', self.cookies.get_cookie_header()
 
		if res.status >= 300 and res.status <= 399:
			location = urlparse.urlparse(res.getheader('Location'))
			if location.netloc != self.host: 
				raise errors.HTTPError, 'Redirecting to different hosts not supported'
			if res.status in (302, 303):
				del headers['Content-Type']
				method = 'GET'
				data = ''
			path = location.path
			if location.query: path = path + '?' + location.query
			res.read()
			return self.__request(method, path, headers, data)
 
		if res.status != 200:
			raise errors.HTTPError, (res.status, res)
 
		return res
 
	def get(self, path, headers = None, data = None):
		return self.__request('GET', path, headers, data)
	def post(self, path, headers = None, data = None):
		return self.__request('POST', path, headers, data)
	def head(self, path, headers = None):
		res = self.__request('HEAD', path, headers, None)
		res.read()
		return res.status, res.getheaders()