Increasing super_len compatibilty to include BytesIO and cStringIO objects.

Added a check for 'getvalue' attr, calling it to retrieve the length if we can.

We also try/except the fileno() call, which can throw
io.UnsupportedOperation for BytesIO because, well, they're not files.
This commit is contained in:
Matt Spitz
2013-10-10 14:54:47 -04:00
parent 025543b604
commit 13a6e02ccd
2 changed files with 34 additions and 2 deletions
+13 -2
View File
@@ -12,6 +12,7 @@ that are also useful for external consumption.
import cgi
import codecs
import collections
import io
import os
import platform
import re
@@ -46,11 +47,21 @@ def dict_to_sequence(d):
def super_len(o):
if hasattr(o, '__len__'):
return len(o)
if hasattr(o, 'len'):
return o.len
if hasattr(o, 'fileno'):
return os.fstat(o.fileno()).st_size
if hasattr(o, 'fileno'):
try:
fileno = o.fileno()
except io.UnsupportedOperation:
pass
else:
return os.fstat(fileno).st_size
if hasattr(o, 'getvalue'):
# e.g. BytesIO, cStringIO.StringI
return len(o.getvalue())
def get_netrc_auth(url):
"""Returns the Requests tuple auth for a given url from netrc."""
+21
View File
@@ -873,5 +873,26 @@ class TestCaseInsensitiveDict(unittest.TestCase):
self.assertEqual(frozenset(cid), keyset)
class UtilsTestCase(unittest.TestCase):
def test_super_len_io_streams(self):
""" Ensures that we properly deal with different kinds of IO streams. """
# uses StringIO or io.StringIO (see import above)
from io import BytesIO
from requests.utils import super_len
self.assertEqual(super_len(StringIO.StringIO()), 0)
self.assertEqual(super_len(StringIO.StringIO('with so much drama in the LBC')), 29)
self.assertEqual(super_len(BytesIO()), 0)
self.assertEqual(super_len(BytesIO(b"it's kinda hard bein' snoop d-o-double-g")), 40)
try:
import cStringIO
except ImportError:
pass
else:
self.assertEqual(super_len(cStringIO.StringIO('but some how, some way...')), 25)
if __name__ == '__main__':
unittest.main()