Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Refactor client response decoding #98

Merged
merged 5 commits into from
Jul 3, 2014
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions README.rst
Original file line number Diff line number Diff line change
Expand Up @@ -29,13 +29,13 @@ To retrieve something from the web::

def get_body(url):
response = yield from aiohttp.request('GET', url)
return (yield from response.read_and_close())
return (yield from response.read())

You can use the get command like this anywhere in your ``asyncio``
powered program::

response = yield from aiohttp.request('GET', 'http://python.org')
body = yield from response.read_and_close()
body = yield from response.read()
print(body)

The signature of request is the following::
Expand Down
50 changes: 37 additions & 13 deletions aiohttp/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@
import warnings

import aiohttp
from .helpers import parse_mimetype
from .log import client_log
from .multidict import CaseInsensitiveMultiDict, MultiDict, MutableMultiDict

Expand Down Expand Up @@ -86,7 +87,7 @@ def request(method, url, *,
>>> resp = yield from aiohttp.request('GET', 'http://python.org/')
>>> resp
<ClientResponse(python.org/) [200]>
>>> data = yield from resp.read_and_close()
>>> data = yield from resp.read()

"""
redirects = 0
Expand Down Expand Up @@ -674,6 +675,10 @@ def close(self, force=False):
self._writer = None
self._writer_wr = None

@asyncio.coroutine
def release(self):
yield from self.read()

@asyncio.coroutine
def wait_for_close(self):
if self._writer is not None:
Expand All @@ -697,7 +702,10 @@ def read(self, decode=False):
buf.append((chunk, size))
total += size
except aiohttp.EofStream:
pass
self.close()
except:
self.close(True)
raise

self._content = bytearray(total)

Expand All @@ -710,24 +718,40 @@ def read(self, decode=False):
data = self._content

if decode:
ct = self.headers.get('CONTENT-TYPE', '').lower()
if ct == 'application/json':
data = json.loads(data.decode('utf-8'))
warnings.warn(
'.read(True) is deprecated. use .json() instead',
DeprecationWarning
)
return (yield from self.json())

return data

@asyncio.coroutine
def read_and_close(self, decode=False):
"""Read response payload and then close response."""
try:
payload = yield from self.read(decode)
except:
self.close(True)
raise
else:
self.close()
warnings.warn(
'read_and_close is deprecated, use .read() instead',
DeprecationWarning
)
return (yield from self.read(decode))

@asyncio.coroutine
def json(self, *, encoding=None):
"""Reads and decodes JSON response."""
if self._content is None:
yield from self.read()

ctype = self.headers.get('CONTENT-TYPE', '').lower()
mtype, stype, _, params = parse_mimetype(ctype)
if not (mtype == 'application' or stype == 'json'):
client_log.warning(
'Attempt to decode JSON with unexpected mimetype: %s', ctype)

if not self._content.strip():
return None

return payload
encoding = encoding or params.get('charset', 'utf-8')
return json.loads(self._content.decode(encoding))


def str_to_bytes(s, encoding='utf-8'):
Expand Down
53 changes: 53 additions & 0 deletions aiohttp/helpers.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,53 @@
"""Various helper functions"""


def parse_mimetype(mimetype):
"""Parses a MIME type into it components.

:param str mimetype: MIME type

:returns: 4 element tuple for MIME type, subtype, suffix and parameters
:rtype: tuple

>>> parse_mimetype('*')
('*', '*', '', {})

>>> parse_mimetype('application/json')
('application', 'json', '', {})

>>> parse_mimetype('application/json; charset=utf-8')
('application', 'json', '', {'charset': 'utf-8'})

>>> parse_mimetype('''application/json;
... charset=utf-8;''')
('application', 'json', '', {'charset': 'utf-8'})

>>> parse_mimetype('ApPlIcAtIoN/JSON;ChaRseT="UTF-8"')
('application', 'json', '', {'charset': 'UTF-8'})

>>> parse_mimetype('application/rss+xml')
('application', 'rss', 'xml', {})

>>> parse_mimetype('text/plain;base64')
('text', 'plain', '', {'base64': ''})

"""
if not mimetype:
return '', '', '', {}

parts = mimetype.split(';')
params = []
for item in parts[1:]:
if not item:
continue
key, value = item.split('=', 2) if '=' in item else (item, '')
params.append((key.lower().strip(), value.strip(' "')))
params = dict(params)

fulltype = parts[0].strip().lower()
if fulltype == '*':
fulltype = '*/*'
mtype, stype = fulltype.split('/', 2) if '/' in fulltype else (fulltype, '')
stype, suffix = stype.split('+') if '+' in stype else (stype, '')

return mtype, stype, suffix, params
71 changes: 61 additions & 10 deletions tests/test_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -81,28 +81,79 @@ def test_repr(self):
repr(self.response))

def test_read_and_close(self):
self.response.read = unittest.mock.Mock()
self.response.read.return_value = asyncio.Future(loop=self.loop)
self.response.read.return_value.set_result(b'payload')
def side_effect(*args, **kwargs):
def second_call(*args, **kwargs):
raise aiohttp.EofStream
fut = asyncio.Future(loop=self.loop)
fut.set_result(b'payload')
content.read.side_effect = second_call
return fut
content = self.response.content = unittest.mock.Mock()
content.read.side_effect = side_effect
self.response.close = unittest.mock.Mock()

res = self.loop.run_until_complete(self.response.read_and_close())
res = self.loop.run_until_complete(self.response.read())
self.assertEqual(res, b'payload')
self.assertTrue(self.response.read.called)
self.assertTrue(self.response.close.called)

def test_read_and_close_with_error(self):
self.response.read = unittest.mock.Mock()
self.response.read.return_value = asyncio.Future(loop=self.loop)
self.response.read.return_value.set_exception(ValueError)
content = self.response.content = unittest.mock.Mock()
content.read.return_value = asyncio.Future(loop=self.loop)
content.read.return_value.set_exception(ValueError)
self.response.close = unittest.mock.Mock()

self.assertRaises(
ValueError,
self.loop.run_until_complete, self.response.read_and_close())
self.assertTrue(self.response.read.called)
self.loop.run_until_complete, self.response.read())
self.response.close.assert_called_with(True)

def test_release(self):
fut = asyncio.Future(loop=self.loop)
fut.set_exception(aiohttp.EofStream)
content = self.response.content = unittest.mock.Mock()
content.read.return_value = fut
self.response.close = unittest.mock.Mock()

self.loop.run_until_complete(self.response.release())
self.assertTrue(self.response.close.called)

def test_json(self):
def side_effect(*args, **kwargs):
def second_call(*args, **kwargs):
raise aiohttp.EofStream
fut = asyncio.Future(loop=self.loop)
fut.set_result('{"тест": "пройден"}'.encode('cp1251'))
content.read.side_effect = second_call
return fut
self.response.headers = {
'CONTENT-TYPE': 'application/json;charset=cp1251'}
content = self.response.content = unittest.mock.Mock()
content.read.side_effect = side_effect
self.response.close = unittest.mock.Mock()

res = self.loop.run_until_complete(self.response.json())
self.assertEqual(res, {'тест': 'пройден'})
self.assertTrue(self.response.close.called)

def test_json_override_encoding(self):
def side_effect(*args, **kwargs):
def second_call(*args, **kwargs):
raise aiohttp.EofStream
fut = asyncio.Future(loop=self.loop)
fut.set_result('{"тест": "пройден"}'.encode('cp1251'))
content.read.side_effect = second_call
return fut
self.response.headers = {
'CONTENT-TYPE': 'application/json;charset=utf8'}
content = self.response.content = unittest.mock.Mock()
content.read.side_effect = side_effect
self.response.close = unittest.mock.Mock()

res = self.loop.run_until_complete(
self.response.json(encoding='cp1251'))
self.assertEqual(res, {'тест': 'пройден'})
self.assertTrue(self.response.close.called)


class ClientRequestTests(unittest.TestCase):

Expand Down
Loading