test_ioutils.py 19 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511
  1. import io
  2. import os
  3. import codecs
  4. import random
  5. import string
  6. try:
  7. from StringIO import StringIO
  8. except ImportError:
  9. # py3
  10. StringIO = io.StringIO
  11. from tempfile import mkdtemp
  12. from unittest import TestCase
  13. from zipfile import ZipFile, ZIP_DEFLATED
  14. from boltons import ioutils
  15. CUR_FILE_PATH = os.path.abspath(__file__)
  16. try:
  17. text_type = unicode # Python 2
  18. binary_type = str
  19. except NameError:
  20. text_type = str # Python 3
  21. binary_type = bytes
  22. class AssertionsMixin(object):
  23. def assertIsNone(self, item, msg=None):
  24. self.assertTrue(item is None, msg)
  25. class BaseTestMixin(object):
  26. """
  27. A set of tests that work the same for SpooledBtyesIO and SpooledStringIO
  28. """
  29. def test_getvalue_norollover(self):
  30. """Make sure getvalue function works with in-memory flo"""
  31. self.spooled_flo.write(self.test_str)
  32. self.assertEqual(self.spooled_flo.getvalue(), self.test_str)
  33. def test_getvalue_rollover(self):
  34. """Make sure getvalue function works with on-disk flo"""
  35. self.spooled_flo.write(self.test_str)
  36. self.assertFalse(self.spooled_flo._rolled)
  37. self.spooled_flo.rollover()
  38. self.assertEqual(self.spooled_flo.getvalue(), self.test_str)
  39. self.assertTrue(self.spooled_flo._rolled)
  40. def test_rollover_custom_directory(self):
  41. """dir keyword argument is passed to TemporaryFile instantiation"""
  42. custom_dir = mkdtemp()
  43. try:
  44. # Re-instantiate self.spooled_flo with the custom dir argument
  45. _spooled_flo = type(self.spooled_flo)(dir=custom_dir)
  46. self.assertEqual(_spooled_flo._dir, custom_dir)
  47. # TemporaryFile is kind of a black box, we can't really test it
  48. # since the directory entry for the file is removed immediately
  49. # after the file is created. So we can't check path using fd.name
  50. # or listdir(custom_dir). We could either convert rollover() to
  51. # use NamedtemporaryFile-s or assume it's well tested enough that
  52. # passing dir= into the constructor will work as expected. We'll
  53. # call rollover() with the dir attribute set just to ensure
  54. # nothing has gone absurdly wrong.
  55. _spooled_flo.write(self.test_str)
  56. _spooled_flo.rollover()
  57. self.assertEqual(_spooled_flo.getvalue(), self.test_str)
  58. self.assertTrue(_spooled_flo._rolled)
  59. _spooled_flo.close()
  60. finally:
  61. os.rmdir(custom_dir)
  62. def test_compare_err(self):
  63. """Read-heads are reset if a comparison raises an error."""
  64. def _monkey_err(*args, **kwargs):
  65. raise Exception('A sad error has occurred today')
  66. a = self.spooled_flo.__class__()
  67. a.write(self.test_str)
  68. b = self.spooled_flo.__class__()
  69. b.write(self.test_str)
  70. a.seek(1)
  71. b.seek(2)
  72. b.__next__ = _monkey_err
  73. try:
  74. a == b
  75. except Exception:
  76. pass
  77. self.assertEqual(a.tell(), 1)
  78. self.assertEqual(b.tell(), 2)
  79. def test_truncate_noargs_norollover(self):
  80. """Test truncating with no args with in-memory flo"""
  81. self.spooled_flo.write(self.test_str)
  82. self.spooled_flo.seek(10)
  83. self.spooled_flo.truncate()
  84. self.assertEqual(self.spooled_flo.getvalue(), self.test_str[:10])
  85. def test_truncate_noargs_rollover(self):
  86. """Test truncating with no args with on-disk flo"""
  87. self.spooled_flo.write(self.test_str)
  88. self.spooled_flo.seek(10)
  89. self.spooled_flo.rollover()
  90. self.spooled_flo.truncate()
  91. self.assertEqual(self.spooled_flo.getvalue(), self.test_str[:10])
  92. def test_truncate_with_args_norollover(self):
  93. """Test truncating to a value with in-memory flo"""
  94. self.spooled_flo.write(self.test_str)
  95. self.spooled_flo.seek(5)
  96. self.spooled_flo.truncate(10)
  97. self.assertEqual(self.spooled_flo.getvalue(), self.test_str[:10])
  98. def test_truncate_with_args_rollover(self):
  99. """Test truncating to a value with on-disk flo"""
  100. self.spooled_flo.write(self.test_str)
  101. self.spooled_flo.seek(5)
  102. self.spooled_flo.rollover()
  103. self.spooled_flo.truncate(10)
  104. self.assertEqual(self.spooled_flo.getvalue(), self.test_str[:10])
  105. def test_type_error_too_many_args(self):
  106. """Make sure TypeError raised if too many args passed to truncate"""
  107. self.spooled_flo.write(self.test_str)
  108. self.assertRaises(TypeError, self.spooled_flo.truncate, 0, 10)
  109. def test_io_error_negative_truncate(self):
  110. """Make sure IOError raised trying to truncate with negative value"""
  111. self.spooled_flo.write(self.test_str)
  112. self.assertRaises(IOError, self.spooled_flo.truncate, -1)
  113. def test_compare_different_instances(self):
  114. """Make sure two different instance types are not considered equal"""
  115. a = ioutils.SpooledBytesIO()
  116. a.write(binary_type(b"I am equal!"))
  117. b = ioutils.SpooledStringIO()
  118. b.write(text_type("I am equal!"))
  119. self.assertNotEqual(a, b)
  120. def test_compare_unequal_instances(self):
  121. """Comparisons of non-SpooledIOBase classes should fail"""
  122. self.assertNotEqual("Bummer dude", self.spooled_flo)
  123. def test_set_softspace_attribute(self):
  124. """Ensure softspace attribute can be retrieved and set"""
  125. self.spooled_flo.softspace = True
  126. self.assertTrue(self.spooled_flo.softspace)
  127. def test_set_softspace_attribute_rolled(self):
  128. """Ensure softspace attribute can be retrieved and set if rolled"""
  129. self.spooled_flo.softspace = True
  130. self.assertTrue(self.spooled_flo.softspace)
  131. self.spooled_flo.rollover()
  132. self.spooled_flo.softspace = True
  133. self.assertTrue(self.spooled_flo.softspace)
  134. def test_buf_property(self):
  135. """'buf' property returns the same value as getvalue()"""
  136. self.assertEqual(self.spooled_flo.buf, self.spooled_flo.getvalue())
  137. def test_pos_property(self):
  138. """'pos' property returns the same value as tell()"""
  139. self.assertEqual(self.spooled_flo.pos, self.spooled_flo.tell())
  140. def test_closed_property(self):
  141. """'closed' property works as expected"""
  142. self.assertFalse(self.spooled_flo.closed)
  143. self.spooled_flo.close()
  144. self.assertTrue(self.spooled_flo.closed)
  145. def test_readline(self):
  146. """Make readline returns expected values"""
  147. self.spooled_flo.write(self.test_str_lines)
  148. self.spooled_flo.seek(0)
  149. self.assertEqual(self.spooled_flo.readline().rstrip(self.linesep),
  150. self.test_str_lines.split(self.linesep)[0])
  151. def test_readlines(self):
  152. """Make sure readlines returns expected values"""
  153. self.spooled_flo.write(self.test_str_lines)
  154. self.spooled_flo.seek(0)
  155. self.assertEqual(
  156. [x.rstrip(self.linesep) for x in self.spooled_flo.readlines()],
  157. self.test_str_lines.split(self.linesep)
  158. )
  159. def test_next(self):
  160. """Make next returns expected values"""
  161. self.spooled_flo.write(self.test_str_lines)
  162. self.spooled_flo.seek(0)
  163. self.assertEqual(self.spooled_flo.next().rstrip(self.linesep),
  164. self.test_str_lines.split(self.linesep)[0])
  165. def test_isatty(self):
  166. """Make sure we can check if the value is a tty"""
  167. # This should simply not fail
  168. self.assertTrue(self.spooled_flo.isatty() is True or
  169. self.spooled_flo.isatty() is False)
  170. def test_truthy(self):
  171. """Make sure empty instances are still considered truthy"""
  172. self.spooled_flo.seek(0)
  173. self.spooled_flo.truncate()
  174. if not self.spooled_flo:
  175. raise AssertionError("Instance is not truthy")
  176. def test_instance_check(self):
  177. """Instance checks against IOBase succeed."""
  178. if not isinstance(self.spooled_flo, io.IOBase):
  179. raise AssertionError('{} is not an instance of IOBase'.format(type(self.spooled_flo)))
  180. def test_closed_file_method_valueerrors(self):
  181. """ValueError raised on closed files for certain methods."""
  182. self.spooled_flo.close()
  183. methods = (
  184. 'flush', 'isatty', 'pos', 'buf', 'truncate', '__next__', '__iter__',
  185. '__enter__', 'read', 'readline', 'tell',
  186. )
  187. for method_name in methods:
  188. with self.assertRaises(ValueError):
  189. getattr(self.spooled_flo, method_name)()
  190. class TestSpooledBytesIO(TestCase, BaseTestMixin, AssertionsMixin):
  191. linesep = os.linesep.encode('ascii')
  192. def setUp(self):
  193. self.spooled_flo = ioutils.SpooledBytesIO()
  194. self.test_str = b"Armado en los EE, UU. para S. P. Richards co.,"
  195. self.test_str_lines = (
  196. "Text with:{0}newlines!".format(os.linesep).encode('ascii')
  197. )
  198. self.data_type = binary_type
  199. def test_compare_not_equal_instances(self):
  200. """Make sure instances with different values fail == check."""
  201. a = ioutils.SpooledBytesIO()
  202. a.write(b"I am a!")
  203. b = ioutils.SpooledBytesIO()
  204. b.write(b"I am b!")
  205. self.assertNotEqual(a, b)
  206. def test_compare_two_equal_instances(self):
  207. """Make sure we can compare instances"""
  208. a = ioutils.SpooledBytesIO()
  209. a.write(b"I am equal!")
  210. b = ioutils.SpooledBytesIO()
  211. b.write(b"I am equal!")
  212. self.assertEqual(a, b)
  213. def test_auto_rollover(self):
  214. """Make sure file rolls over to disk after max_size reached"""
  215. tmp = ioutils.SpooledBytesIO(max_size=10)
  216. tmp.write(b"The quick brown fox jumped over the lazy dogs.")
  217. self.assertTrue(tmp._rolled)
  218. def test_use_as_context_mgr(self):
  219. """Make sure SpooledBytesIO can be used as a context manager"""
  220. test_str = b"Armado en los EE, UU. para S. P. Richards co.,"
  221. with ioutils.SpooledBytesIO() as f:
  222. f.write(test_str)
  223. self.assertEqual(f.getvalue(), test_str)
  224. def test_len_no_rollover(self):
  225. """Make sure len works with in-memory flo"""
  226. self.spooled_flo.write(self.test_str)
  227. self.assertEqual(self.spooled_flo.len, len(self.test_str))
  228. self.assertEqual(len(self.spooled_flo), len(self.test_str))
  229. def test_len_rollover(self):
  230. """Make sure len works with on-disk flo"""
  231. self.spooled_flo.write(self.test_str)
  232. self.spooled_flo.rollover()
  233. self.assertEqual(self.spooled_flo.len, len(self.test_str))
  234. self.assertEqual(len(self.spooled_flo), len(self.test_str))
  235. def test_invalid_type(self):
  236. """Ensure TypeError raised when writing unicode to SpooledBytesIO"""
  237. self.assertRaises(TypeError, self.spooled_flo.write, u"hi")
  238. def test_flush_after_rollover(self):
  239. """Make sure we can flush before and after rolling to a real file"""
  240. self.spooled_flo.write(self.test_str)
  241. self.assertIsNone(self.spooled_flo.flush())
  242. self.spooled_flo.rollover()
  243. self.assertIsNone(self.spooled_flo.flush())
  244. def test_zip_compat(self):
  245. """Make sure object is compatible with ZipFile library"""
  246. self.spooled_flo.seek(0)
  247. self.spooled_flo.truncate()
  248. doc = ZipFile(self.spooled_flo, 'w', ZIP_DEFLATED)
  249. doc.writestr("content.txt", "test")
  250. self.assertTrue('content.txt' in doc.namelist())
  251. doc.close()
  252. def test_iter(self):
  253. """Make sure iter works as expected"""
  254. self.spooled_flo.write(b"a\nb")
  255. self.spooled_flo.seek(0)
  256. self.assertEqual([x for x in self.spooled_flo], [b"a\n", b"b"])
  257. def test_writelines(self):
  258. """An iterable of lines can be written"""
  259. lines = [b"1", b"2", b"3"]
  260. expected = b"123"
  261. self.spooled_flo.writelines(lines)
  262. self.assertEqual(self.spooled_flo.getvalue(), expected)
  263. class TestSpooledStringIO(TestCase, BaseTestMixin, AssertionsMixin):
  264. linesep = os.linesep
  265. def setUp(self):
  266. self.spooled_flo = ioutils.SpooledStringIO()
  267. self.test_str = u"Remember kids, always use an emdash: '\u2014'"
  268. self.test_str_lines = u"Text with\u2014{0}newlines!".format(os.linesep)
  269. self.data_type = text_type
  270. def test_compare_not_equal_instances(self):
  271. """Make sure instances with different values fail == check."""
  272. a = ioutils.SpooledStringIO()
  273. a.write(u"I am a!")
  274. b = ioutils.SpooledStringIO()
  275. b.write(u"I am b!")
  276. self.assertNotEqual(a, b)
  277. def test_compare_two_equal_instances(self):
  278. """Make sure we can compare instances"""
  279. a = ioutils.SpooledStringIO()
  280. a.write(u"I am equal!")
  281. b = ioutils.SpooledStringIO()
  282. b.write(u"I am equal!")
  283. self.assertEqual(a, b)
  284. def test_auto_rollover(self):
  285. """Make sure file rolls over to disk after max_size reached"""
  286. tmp = ioutils.SpooledStringIO(max_size=10)
  287. tmp.write(u"The quick brown fox jumped over the lazy dogs.")
  288. self.assertTrue(tmp._rolled)
  289. def test_use_as_context_mgr(self):
  290. """Make sure SpooledStringIO can be used as a context manager"""
  291. test_str = u"Armado en los EE, UU. para S. P. Richards co.,"
  292. with ioutils.SpooledStringIO() as f:
  293. f.write(test_str)
  294. self.assertEqual(f.getvalue(), test_str)
  295. def test_len_no_rollover(self):
  296. """Make sure len property works with in-memory flo"""
  297. self.spooled_flo.write(self.test_str)
  298. self.assertEqual(self.spooled_flo.len, len(self.test_str))
  299. def test_len_rollover(self):
  300. """Make sure len property works with on-disk flo"""
  301. self.spooled_flo.write(self.test_str)
  302. self.spooled_flo.rollover()
  303. self.assertEqual(self.spooled_flo.len, len(self.test_str))
  304. def test_invalid_type(self):
  305. """Ensure TypeError raised when writing bytes to SpooledStringIO"""
  306. self.assertRaises(TypeError, self.spooled_flo.write, b"hi")
  307. def test_tell_codepoints(self):
  308. """Verify tell() returns codepoint position, not bytes position"""
  309. self.spooled_flo.write(self.test_str)
  310. self.spooled_flo.seek(0)
  311. self.spooled_flo.read(40)
  312. self.assertEqual(self.spooled_flo.tell(), 40)
  313. self.spooled_flo.seek(10)
  314. self.assertEqual(self.spooled_flo.tell(), 10)
  315. def test_codepoints_all_enc(self):
  316. """"Test getting read, seek, tell, on various codepoints"""
  317. test_str = u"\u2014\u2014\u2014"
  318. self.spooled_flo.write(test_str)
  319. self.spooled_flo.seek(1)
  320. self.assertEqual(self.spooled_flo.read(), u"\u2014\u2014")
  321. self.assertEqual(len(self.spooled_flo), len(test_str))
  322. def test_seek_codepoints_SEEK_END(self):
  323. """Make sure seek() moves to codepoints relative to file end"""
  324. self.spooled_flo.write(self.test_str)
  325. ret = self.spooled_flo.seek(0, os.SEEK_END)
  326. self.assertEqual(ret, len(self.test_str))
  327. def test_seek_codepoints_large_SEEK_END(self):
  328. """Make sure seek() moves to codepoints relative to file end"""
  329. test_str = u"".join(random.choice(string.ascii_letters) for
  330. x in range(34000))
  331. self.spooled_flo.write(test_str)
  332. ret = self.spooled_flo.seek(0, os.SEEK_END)
  333. self.assertEqual(ret, len(test_str))
  334. def test_seek_codepoints_SEEK_SET(self):
  335. """Make sure seek() moves to codepoints relative to file start"""
  336. self.spooled_flo.write(self.test_str)
  337. ret = self.spooled_flo.seek(3, os.SEEK_SET)
  338. self.assertEqual(ret, 3)
  339. def test_seek_codepoints_large_SEEK_SET(self):
  340. """Make sure seek() moves to codepoints relative to file start"""
  341. test_str = u"".join(random.choice(string.ascii_letters) for
  342. x in range(34000))
  343. self.spooled_flo.write(test_str)
  344. ret = self.spooled_flo.seek(33000, os.SEEK_SET)
  345. self.assertEqual(ret, 33000)
  346. def test_seek_codepoints_SEEK_CUR(self):
  347. """Make sure seek() moves to codepoints relative to current_position"""
  348. test_str = u"\u2014\u2014\u2014"
  349. self.spooled_flo.write(test_str)
  350. self.spooled_flo.seek(1)
  351. self.assertEqual(self.spooled_flo.tell(), 1)
  352. ret = self.spooled_flo.seek(2, os.SEEK_CUR)
  353. self.assertEqual(ret, 3)
  354. def test_seek_codepoints_large_SEEK_CUR(self):
  355. """Make sure seek() moves to codepoints relative to current_position"""
  356. test_str = u"".join(random.choice(string.ascii_letters) for
  357. x in range(34000))
  358. self.spooled_flo.write(test_str)
  359. self.spooled_flo.seek(1)
  360. ret = self.spooled_flo.seek(33000, os.SEEK_CUR)
  361. self.assertEqual(ret, 33001)
  362. def test_x80_codepoint(self):
  363. """Make sure x80 codepoint doesn't confuse read value"""
  364. test_str = u'\x8000'
  365. self.spooled_flo.write(test_str)
  366. self.spooled_flo.seek(0)
  367. self.assertEqual(len(self.spooled_flo.read(2)), 2)
  368. self.assertEqual(self.spooled_flo.read(), '0')
  369. def test_seek_encoded(self):
  370. """Make sure reading works when bytes exceeds read val"""
  371. test_str = u"\u2014\u2014\u2014"
  372. self.spooled_flo.write(test_str)
  373. self.spooled_flo.seek(0)
  374. self.assertEqual(self.spooled_flo.read(3), test_str)
  375. def test_iter(self):
  376. """Make sure iter works as expected"""
  377. self.spooled_flo.write(u"a\nb")
  378. self.spooled_flo.seek(0)
  379. self.assertEqual([x for x in self.spooled_flo], [u"a\n", u"b"])
  380. def test_writelines(self):
  381. """An iterable of lines can be written"""
  382. lines = [u"1", u"2", u"3"]
  383. expected = u"123"
  384. self.spooled_flo.writelines(lines)
  385. self.assertEqual(self.spooled_flo.getvalue(), expected)
  386. class TestMultiFileReader(TestCase):
  387. def test_read_seek_bytes(self):
  388. r = ioutils.MultiFileReader(io.BytesIO(b'narf'), io.BytesIO(b'troz'))
  389. self.assertEqual([b'nar', b'ftr', b'oz'],
  390. list(iter(lambda: r.read(3), b'')))
  391. r.seek(0)
  392. self.assertEqual(b'narftroz', r.read())
  393. def test_read_seek_text(self):
  394. # also tests StringIO.StringIO on py2
  395. r = ioutils.MultiFileReader(StringIO(u'narf'),
  396. io.StringIO(u'troz'))
  397. self.assertEqual([u'nar', u'ftr', u'oz'],
  398. list(iter(lambda: r.read(3), u'')))
  399. r.seek(0)
  400. self.assertEqual(u'narftroz', r.read())
  401. def test_no_mixed_bytes_and_text(self):
  402. self.assertRaises(ValueError, ioutils.MultiFileReader,
  403. io.BytesIO(b'narf'), io.StringIO(u'troz'))
  404. def test_open(self):
  405. with open(CUR_FILE_PATH, 'r') as f:
  406. r_file_str = f.read()
  407. with open(CUR_FILE_PATH, 'r') as f1:
  408. with open(CUR_FILE_PATH, 'r') as f2:
  409. mfr = ioutils.MultiFileReader(f1, f2)
  410. r_double_file_str = mfr.read()
  411. assert r_double_file_str == (r_file_str * 2)
  412. with open(CUR_FILE_PATH, 'rb') as f:
  413. rb_file_str = f.read()
  414. with open(CUR_FILE_PATH, 'rb') as f1:
  415. with open(CUR_FILE_PATH, 'rb') as f2:
  416. mfr = ioutils.MultiFileReader(f1, f2)
  417. rb_double_file_str = mfr.read()
  418. assert rb_double_file_str == (rb_file_str * 2)
  419. utf8_file_str = codecs.open(CUR_FILE_PATH, encoding='utf8').read()
  420. f1, f2 = (codecs.open(CUR_FILE_PATH, encoding='utf8'),
  421. codecs.open(CUR_FILE_PATH, encoding='utf8'))
  422. mfr = ioutils.MultiFileReader(f1, f2)
  423. utf8_double_file_str = mfr.read()
  424. assert utf8_double_file_str == (utf8_file_str * 2)