123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511 |
- import io
- import os
- import codecs
- import random
- import string
- try:
- from StringIO import StringIO
- except ImportError:
- # py3
- StringIO = io.StringIO
- from tempfile import mkdtemp
- from unittest import TestCase
- from zipfile import ZipFile, ZIP_DEFLATED
- from boltons import ioutils
- CUR_FILE_PATH = os.path.abspath(__file__)
- try:
- text_type = unicode # Python 2
- binary_type = str
- except NameError:
- text_type = str # Python 3
- binary_type = bytes
- class AssertionsMixin(object):
- def assertIsNone(self, item, msg=None):
- self.assertTrue(item is None, msg)
- class BaseTestMixin(object):
- """
- A set of tests that work the same for SpooledBtyesIO and SpooledStringIO
- """
- def test_getvalue_norollover(self):
- """Make sure getvalue function works with in-memory flo"""
- self.spooled_flo.write(self.test_str)
- self.assertEqual(self.spooled_flo.getvalue(), self.test_str)
- def test_getvalue_rollover(self):
- """Make sure getvalue function works with on-disk flo"""
- self.spooled_flo.write(self.test_str)
- self.assertFalse(self.spooled_flo._rolled)
- self.spooled_flo.rollover()
- self.assertEqual(self.spooled_flo.getvalue(), self.test_str)
- self.assertTrue(self.spooled_flo._rolled)
- def test_rollover_custom_directory(self):
- """dir keyword argument is passed to TemporaryFile instantiation"""
- custom_dir = mkdtemp()
- try:
- # Re-instantiate self.spooled_flo with the custom dir argument
- _spooled_flo = type(self.spooled_flo)(dir=custom_dir)
- self.assertEqual(_spooled_flo._dir, custom_dir)
- # TemporaryFile is kind of a black box, we can't really test it
- # since the directory entry for the file is removed immediately
- # after the file is created. So we can't check path using fd.name
- # or listdir(custom_dir). We could either convert rollover() to
- # use NamedtemporaryFile-s or assume it's well tested enough that
- # passing dir= into the constructor will work as expected. We'll
- # call rollover() with the dir attribute set just to ensure
- # nothing has gone absurdly wrong.
- _spooled_flo.write(self.test_str)
- _spooled_flo.rollover()
- self.assertEqual(_spooled_flo.getvalue(), self.test_str)
- self.assertTrue(_spooled_flo._rolled)
- _spooled_flo.close()
- finally:
- os.rmdir(custom_dir)
- def test_compare_err(self):
- """Read-heads are reset if a comparison raises an error."""
- def _monkey_err(*args, **kwargs):
- raise Exception('A sad error has occurred today')
- a = self.spooled_flo.__class__()
- a.write(self.test_str)
- b = self.spooled_flo.__class__()
- b.write(self.test_str)
- a.seek(1)
- b.seek(2)
- b.__next__ = _monkey_err
- try:
- a == b
- except Exception:
- pass
- self.assertEqual(a.tell(), 1)
- self.assertEqual(b.tell(), 2)
- def test_truncate_noargs_norollover(self):
- """Test truncating with no args with in-memory flo"""
- self.spooled_flo.write(self.test_str)
- self.spooled_flo.seek(10)
- self.spooled_flo.truncate()
- self.assertEqual(self.spooled_flo.getvalue(), self.test_str[:10])
- def test_truncate_noargs_rollover(self):
- """Test truncating with no args with on-disk flo"""
- self.spooled_flo.write(self.test_str)
- self.spooled_flo.seek(10)
- self.spooled_flo.rollover()
- self.spooled_flo.truncate()
- self.assertEqual(self.spooled_flo.getvalue(), self.test_str[:10])
- def test_truncate_with_args_norollover(self):
- """Test truncating to a value with in-memory flo"""
- self.spooled_flo.write(self.test_str)
- self.spooled_flo.seek(5)
- self.spooled_flo.truncate(10)
- self.assertEqual(self.spooled_flo.getvalue(), self.test_str[:10])
- def test_truncate_with_args_rollover(self):
- """Test truncating to a value with on-disk flo"""
- self.spooled_flo.write(self.test_str)
- self.spooled_flo.seek(5)
- self.spooled_flo.rollover()
- self.spooled_flo.truncate(10)
- self.assertEqual(self.spooled_flo.getvalue(), self.test_str[:10])
- def test_type_error_too_many_args(self):
- """Make sure TypeError raised if too many args passed to truncate"""
- self.spooled_flo.write(self.test_str)
- self.assertRaises(TypeError, self.spooled_flo.truncate, 0, 10)
- def test_io_error_negative_truncate(self):
- """Make sure IOError raised trying to truncate with negative value"""
- self.spooled_flo.write(self.test_str)
- self.assertRaises(IOError, self.spooled_flo.truncate, -1)
- def test_compare_different_instances(self):
- """Make sure two different instance types are not considered equal"""
- a = ioutils.SpooledBytesIO()
- a.write(binary_type(b"I am equal!"))
- b = ioutils.SpooledStringIO()
- b.write(text_type("I am equal!"))
- self.assertNotEqual(a, b)
- def test_compare_unequal_instances(self):
- """Comparisons of non-SpooledIOBase classes should fail"""
- self.assertNotEqual("Bummer dude", self.spooled_flo)
- def test_set_softspace_attribute(self):
- """Ensure softspace attribute can be retrieved and set"""
- self.spooled_flo.softspace = True
- self.assertTrue(self.spooled_flo.softspace)
- def test_set_softspace_attribute_rolled(self):
- """Ensure softspace attribute can be retrieved and set if rolled"""
- self.spooled_flo.softspace = True
- self.assertTrue(self.spooled_flo.softspace)
- self.spooled_flo.rollover()
- self.spooled_flo.softspace = True
- self.assertTrue(self.spooled_flo.softspace)
- def test_buf_property(self):
- """'buf' property returns the same value as getvalue()"""
- self.assertEqual(self.spooled_flo.buf, self.spooled_flo.getvalue())
- def test_pos_property(self):
- """'pos' property returns the same value as tell()"""
- self.assertEqual(self.spooled_flo.pos, self.spooled_flo.tell())
- def test_closed_property(self):
- """'closed' property works as expected"""
- self.assertFalse(self.spooled_flo.closed)
- self.spooled_flo.close()
- self.assertTrue(self.spooled_flo.closed)
- def test_readline(self):
- """Make readline returns expected values"""
- self.spooled_flo.write(self.test_str_lines)
- self.spooled_flo.seek(0)
- self.assertEqual(self.spooled_flo.readline().rstrip(self.linesep),
- self.test_str_lines.split(self.linesep)[0])
- def test_readlines(self):
- """Make sure readlines returns expected values"""
- self.spooled_flo.write(self.test_str_lines)
- self.spooled_flo.seek(0)
- self.assertEqual(
- [x.rstrip(self.linesep) for x in self.spooled_flo.readlines()],
- self.test_str_lines.split(self.linesep)
- )
- def test_next(self):
- """Make next returns expected values"""
- self.spooled_flo.write(self.test_str_lines)
- self.spooled_flo.seek(0)
- self.assertEqual(self.spooled_flo.next().rstrip(self.linesep),
- self.test_str_lines.split(self.linesep)[0])
- def test_isatty(self):
- """Make sure we can check if the value is a tty"""
- # This should simply not fail
- self.assertTrue(self.spooled_flo.isatty() is True or
- self.spooled_flo.isatty() is False)
- def test_truthy(self):
- """Make sure empty instances are still considered truthy"""
- self.spooled_flo.seek(0)
- self.spooled_flo.truncate()
- if not self.spooled_flo:
- raise AssertionError("Instance is not truthy")
- def test_instance_check(self):
- """Instance checks against IOBase succeed."""
- if not isinstance(self.spooled_flo, io.IOBase):
- raise AssertionError('{} is not an instance of IOBase'.format(type(self.spooled_flo)))
- def test_closed_file_method_valueerrors(self):
- """ValueError raised on closed files for certain methods."""
- self.spooled_flo.close()
- methods = (
- 'flush', 'isatty', 'pos', 'buf', 'truncate', '__next__', '__iter__',
- '__enter__', 'read', 'readline', 'tell',
- )
- for method_name in methods:
- with self.assertRaises(ValueError):
- getattr(self.spooled_flo, method_name)()
- class TestSpooledBytesIO(TestCase, BaseTestMixin, AssertionsMixin):
- linesep = os.linesep.encode('ascii')
- def setUp(self):
- self.spooled_flo = ioutils.SpooledBytesIO()
- self.test_str = b"Armado en los EE, UU. para S. P. Richards co.,"
- self.test_str_lines = (
- "Text with:{0}newlines!".format(os.linesep).encode('ascii')
- )
- self.data_type = binary_type
- def test_compare_not_equal_instances(self):
- """Make sure instances with different values fail == check."""
- a = ioutils.SpooledBytesIO()
- a.write(b"I am a!")
- b = ioutils.SpooledBytesIO()
- b.write(b"I am b!")
- self.assertNotEqual(a, b)
- def test_compare_two_equal_instances(self):
- """Make sure we can compare instances"""
- a = ioutils.SpooledBytesIO()
- a.write(b"I am equal!")
- b = ioutils.SpooledBytesIO()
- b.write(b"I am equal!")
- self.assertEqual(a, b)
- def test_auto_rollover(self):
- """Make sure file rolls over to disk after max_size reached"""
- tmp = ioutils.SpooledBytesIO(max_size=10)
- tmp.write(b"The quick brown fox jumped over the lazy dogs.")
- self.assertTrue(tmp._rolled)
- def test_use_as_context_mgr(self):
- """Make sure SpooledBytesIO can be used as a context manager"""
- test_str = b"Armado en los EE, UU. para S. P. Richards co.,"
- with ioutils.SpooledBytesIO() as f:
- f.write(test_str)
- self.assertEqual(f.getvalue(), test_str)
- def test_len_no_rollover(self):
- """Make sure len works with in-memory flo"""
- self.spooled_flo.write(self.test_str)
- self.assertEqual(self.spooled_flo.len, len(self.test_str))
- self.assertEqual(len(self.spooled_flo), len(self.test_str))
- def test_len_rollover(self):
- """Make sure len works with on-disk flo"""
- self.spooled_flo.write(self.test_str)
- self.spooled_flo.rollover()
- self.assertEqual(self.spooled_flo.len, len(self.test_str))
- self.assertEqual(len(self.spooled_flo), len(self.test_str))
- def test_invalid_type(self):
- """Ensure TypeError raised when writing unicode to SpooledBytesIO"""
- self.assertRaises(TypeError, self.spooled_flo.write, u"hi")
- def test_flush_after_rollover(self):
- """Make sure we can flush before and after rolling to a real file"""
- self.spooled_flo.write(self.test_str)
- self.assertIsNone(self.spooled_flo.flush())
- self.spooled_flo.rollover()
- self.assertIsNone(self.spooled_flo.flush())
- def test_zip_compat(self):
- """Make sure object is compatible with ZipFile library"""
- self.spooled_flo.seek(0)
- self.spooled_flo.truncate()
- doc = ZipFile(self.spooled_flo, 'w', ZIP_DEFLATED)
- doc.writestr("content.txt", "test")
- self.assertTrue('content.txt' in doc.namelist())
- doc.close()
- def test_iter(self):
- """Make sure iter works as expected"""
- self.spooled_flo.write(b"a\nb")
- self.spooled_flo.seek(0)
- self.assertEqual([x for x in self.spooled_flo], [b"a\n", b"b"])
- def test_writelines(self):
- """An iterable of lines can be written"""
- lines = [b"1", b"2", b"3"]
- expected = b"123"
- self.spooled_flo.writelines(lines)
- self.assertEqual(self.spooled_flo.getvalue(), expected)
- class TestSpooledStringIO(TestCase, BaseTestMixin, AssertionsMixin):
- linesep = os.linesep
- def setUp(self):
- self.spooled_flo = ioutils.SpooledStringIO()
- self.test_str = u"Remember kids, always use an emdash: '\u2014'"
- self.test_str_lines = u"Text with\u2014{0}newlines!".format(os.linesep)
- self.data_type = text_type
- def test_compare_not_equal_instances(self):
- """Make sure instances with different values fail == check."""
- a = ioutils.SpooledStringIO()
- a.write(u"I am a!")
- b = ioutils.SpooledStringIO()
- b.write(u"I am b!")
- self.assertNotEqual(a, b)
- def test_compare_two_equal_instances(self):
- """Make sure we can compare instances"""
- a = ioutils.SpooledStringIO()
- a.write(u"I am equal!")
- b = ioutils.SpooledStringIO()
- b.write(u"I am equal!")
- self.assertEqual(a, b)
- def test_auto_rollover(self):
- """Make sure file rolls over to disk after max_size reached"""
- tmp = ioutils.SpooledStringIO(max_size=10)
- tmp.write(u"The quick brown fox jumped over the lazy dogs.")
- self.assertTrue(tmp._rolled)
- def test_use_as_context_mgr(self):
- """Make sure SpooledStringIO can be used as a context manager"""
- test_str = u"Armado en los EE, UU. para S. P. Richards co.,"
- with ioutils.SpooledStringIO() as f:
- f.write(test_str)
- self.assertEqual(f.getvalue(), test_str)
- def test_len_no_rollover(self):
- """Make sure len property works with in-memory flo"""
- self.spooled_flo.write(self.test_str)
- self.assertEqual(self.spooled_flo.len, len(self.test_str))
- def test_len_rollover(self):
- """Make sure len property works with on-disk flo"""
- self.spooled_flo.write(self.test_str)
- self.spooled_flo.rollover()
- self.assertEqual(self.spooled_flo.len, len(self.test_str))
- def test_invalid_type(self):
- """Ensure TypeError raised when writing bytes to SpooledStringIO"""
- self.assertRaises(TypeError, self.spooled_flo.write, b"hi")
- def test_tell_codepoints(self):
- """Verify tell() returns codepoint position, not bytes position"""
- self.spooled_flo.write(self.test_str)
- self.spooled_flo.seek(0)
- self.spooled_flo.read(40)
- self.assertEqual(self.spooled_flo.tell(), 40)
- self.spooled_flo.seek(10)
- self.assertEqual(self.spooled_flo.tell(), 10)
- def test_codepoints_all_enc(self):
- """"Test getting read, seek, tell, on various codepoints"""
- test_str = u"\u2014\u2014\u2014"
- self.spooled_flo.write(test_str)
- self.spooled_flo.seek(1)
- self.assertEqual(self.spooled_flo.read(), u"\u2014\u2014")
- self.assertEqual(len(self.spooled_flo), len(test_str))
- def test_seek_codepoints_SEEK_END(self):
- """Make sure seek() moves to codepoints relative to file end"""
- self.spooled_flo.write(self.test_str)
- ret = self.spooled_flo.seek(0, os.SEEK_END)
- self.assertEqual(ret, len(self.test_str))
- def test_seek_codepoints_large_SEEK_END(self):
- """Make sure seek() moves to codepoints relative to file end"""
- test_str = u"".join(random.choice(string.ascii_letters) for
- x in range(34000))
- self.spooled_flo.write(test_str)
- ret = self.spooled_flo.seek(0, os.SEEK_END)
- self.assertEqual(ret, len(test_str))
- def test_seek_codepoints_SEEK_SET(self):
- """Make sure seek() moves to codepoints relative to file start"""
- self.spooled_flo.write(self.test_str)
- ret = self.spooled_flo.seek(3, os.SEEK_SET)
- self.assertEqual(ret, 3)
- def test_seek_codepoints_large_SEEK_SET(self):
- """Make sure seek() moves to codepoints relative to file start"""
- test_str = u"".join(random.choice(string.ascii_letters) for
- x in range(34000))
- self.spooled_flo.write(test_str)
- ret = self.spooled_flo.seek(33000, os.SEEK_SET)
- self.assertEqual(ret, 33000)
- def test_seek_codepoints_SEEK_CUR(self):
- """Make sure seek() moves to codepoints relative to current_position"""
- test_str = u"\u2014\u2014\u2014"
- self.spooled_flo.write(test_str)
- self.spooled_flo.seek(1)
- self.assertEqual(self.spooled_flo.tell(), 1)
- ret = self.spooled_flo.seek(2, os.SEEK_CUR)
- self.assertEqual(ret, 3)
- def test_seek_codepoints_large_SEEK_CUR(self):
- """Make sure seek() moves to codepoints relative to current_position"""
- test_str = u"".join(random.choice(string.ascii_letters) for
- x in range(34000))
- self.spooled_flo.write(test_str)
- self.spooled_flo.seek(1)
- ret = self.spooled_flo.seek(33000, os.SEEK_CUR)
- self.assertEqual(ret, 33001)
- def test_x80_codepoint(self):
- """Make sure x80 codepoint doesn't confuse read value"""
- test_str = u'\x8000'
- self.spooled_flo.write(test_str)
- self.spooled_flo.seek(0)
- self.assertEqual(len(self.spooled_flo.read(2)), 2)
- self.assertEqual(self.spooled_flo.read(), '0')
- def test_seek_encoded(self):
- """Make sure reading works when bytes exceeds read val"""
- test_str = u"\u2014\u2014\u2014"
- self.spooled_flo.write(test_str)
- self.spooled_flo.seek(0)
- self.assertEqual(self.spooled_flo.read(3), test_str)
- def test_iter(self):
- """Make sure iter works as expected"""
- self.spooled_flo.write(u"a\nb")
- self.spooled_flo.seek(0)
- self.assertEqual([x for x in self.spooled_flo], [u"a\n", u"b"])
- def test_writelines(self):
- """An iterable of lines can be written"""
- lines = [u"1", u"2", u"3"]
- expected = u"123"
- self.spooled_flo.writelines(lines)
- self.assertEqual(self.spooled_flo.getvalue(), expected)
- class TestMultiFileReader(TestCase):
- def test_read_seek_bytes(self):
- r = ioutils.MultiFileReader(io.BytesIO(b'narf'), io.BytesIO(b'troz'))
- self.assertEqual([b'nar', b'ftr', b'oz'],
- list(iter(lambda: r.read(3), b'')))
- r.seek(0)
- self.assertEqual(b'narftroz', r.read())
- def test_read_seek_text(self):
- # also tests StringIO.StringIO on py2
- r = ioutils.MultiFileReader(StringIO(u'narf'),
- io.StringIO(u'troz'))
- self.assertEqual([u'nar', u'ftr', u'oz'],
- list(iter(lambda: r.read(3), u'')))
- r.seek(0)
- self.assertEqual(u'narftroz', r.read())
- def test_no_mixed_bytes_and_text(self):
- self.assertRaises(ValueError, ioutils.MultiFileReader,
- io.BytesIO(b'narf'), io.StringIO(u'troz'))
- def test_open(self):
- with open(CUR_FILE_PATH, 'r') as f:
- r_file_str = f.read()
- with open(CUR_FILE_PATH, 'r') as f1:
- with open(CUR_FILE_PATH, 'r') as f2:
- mfr = ioutils.MultiFileReader(f1, f2)
- r_double_file_str = mfr.read()
- assert r_double_file_str == (r_file_str * 2)
- with open(CUR_FILE_PATH, 'rb') as f:
- rb_file_str = f.read()
- with open(CUR_FILE_PATH, 'rb') as f1:
- with open(CUR_FILE_PATH, 'rb') as f2:
- mfr = ioutils.MultiFileReader(f1, f2)
- rb_double_file_str = mfr.read()
- assert rb_double_file_str == (rb_file_str * 2)
- utf8_file_str = codecs.open(CUR_FILE_PATH, encoding='utf8').read()
- f1, f2 = (codecs.open(CUR_FILE_PATH, encoding='utf8'),
- codecs.open(CUR_FILE_PATH, encoding='utf8'))
- mfr = ioutils.MultiFileReader(f1, f2)
- utf8_double_file_str = mfr.read()
- assert utf8_double_file_str == (utf8_file_str * 2)
|