1# Tests StringIO and cStringIO
2
3import unittest
4import StringIO
5import cStringIO
6import types
7import array
8import sys
9from test import test_support
10
11
12class TestGenericStringIO(unittest.TestCase):
13    # use a class variable MODULE to define which module is being tested
14
15    # Line of data to test as string
16    _line = 'abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ!'
17
18    # Constructor to use for the test data (._line is passed to this
19    # constructor)
20    constructor = str
21
22    def setUp(self):
23        self._lines = self.constructor((self._line + '\n') * 5)
24        self._fp = self.MODULE.StringIO(self._lines)
25
26    def test_reads(self):
27        eq = self.assertEqual
28        self.assertRaises(TypeError, self._fp.seek)
29        eq(self._fp.read(10), self._line[:10])
30        eq(self._fp.read(0), '')
31        eq(self._fp.readline(0), '')
32        eq(self._fp.readline(), self._line[10:] + '\n')
33        eq(len(self._fp.readlines(60)), 2)
34        self._fp.seek(0)
35        eq(self._fp.readline(-1), self._line + '\n')
36
37    def test_writes(self):
38        f = self.MODULE.StringIO()
39        self.assertRaises(TypeError, f.seek)
40        f.write(self._line[:6])
41        f.seek(3)
42        f.write(self._line[20:26])
43        f.write(self._line[52])
44        self.assertEqual(f.getvalue(), 'abcuvwxyz!')
45
46    def test_writelines(self):
47        f = self.MODULE.StringIO()
48        f.writelines([self._line[0], self._line[1], self._line[2]])
49        f.seek(0)
50        self.assertEqual(f.getvalue(), 'abc')
51
52    def test_writelines_error(self):
53        def errorGen():
54            yield 'a'
55            raise KeyboardInterrupt()
56        f = self.MODULE.StringIO()
57        self.assertRaises(KeyboardInterrupt, f.writelines, errorGen())
58
59    def test_truncate(self):
60        eq = self.assertEqual
61        f = self.MODULE.StringIO()
62        f.write(self._lines)
63        f.seek(10)
64        f.truncate()
65        eq(f.getvalue(), 'abcdefghij')
66        f.truncate(5)
67        eq(f.getvalue(), 'abcde')
68        f.write('xyz')
69        eq(f.getvalue(), 'abcdexyz')
70        self.assertRaises(IOError, f.truncate, -1)
71        f.close()
72        self.assertRaises(ValueError, f.write, 'frobnitz')
73
74    def test_closed_flag(self):
75        f = self.MODULE.StringIO()
76        self.assertEqual(f.closed, False)
77        f.close()
78        self.assertEqual(f.closed, True)
79        f = self.MODULE.StringIO("abc")
80        self.assertEqual(f.closed, False)
81        f.close()
82        self.assertEqual(f.closed, True)
83
84    def test_isatty(self):
85        f = self.MODULE.StringIO()
86        self.assertRaises(TypeError, f.isatty, None)
87        self.assertEqual(f.isatty(), False)
88        f.close()
89        self.assertRaises(ValueError, f.isatty)
90
91    def test_iterator(self):
92        eq = self.assertEqual
93        unless = self.assertTrue
94        eq(iter(self._fp), self._fp)
95        # Does this object support the iteration protocol?
96        unless(hasattr(self._fp, '__iter__'))
97        unless(hasattr(self._fp, 'next'))
98        i = 0
99        for line in self._fp:
100            eq(line, self._line + '\n')
101            i += 1
102        eq(i, 5)
103        self._fp.close()
104        self.assertRaises(ValueError, self._fp.next)
105
106    def test_getvalue(self):
107        self._fp.close()
108        self.assertRaises(ValueError, self._fp.getvalue)
109
110    @test_support.bigmemtest(test_support._2G + 2**26, memuse=2.001)
111    def test_reads_from_large_stream(self, size):
112        linesize = 2**26 # 64 MiB
113        lines = ['x' * (linesize - 1) + '\n'] * (size // linesize) + \
114                ['y' * (size % linesize)]
115        f = self.MODULE.StringIO(''.join(lines))
116        for i, expected in enumerate(lines):
117            line = f.read(len(expected))
118            self.assertEqual(len(line), len(expected))
119            self.assertEqual(line, expected)
120        self.assertEqual(f.read(), '')
121        f.seek(0)
122        for i, expected in enumerate(lines):
123            line = f.readline()
124            self.assertEqual(len(line), len(expected))
125            self.assertEqual(line, expected)
126        self.assertEqual(f.readline(), '')
127        f.seek(0)
128        self.assertEqual(f.readlines(), lines)
129        self.assertEqual(f.readlines(), [])
130        f.seek(0)
131        self.assertEqual(f.readlines(size), lines)
132        self.assertEqual(f.readlines(), [])
133
134    # In worst case cStringIO requires 2 + 1 + 1/2 + 1/2**2 + ... = 4
135    # bytes per input character.
136    @test_support.bigmemtest(test_support._2G, memuse=4)
137    def test_writes_to_large_stream(self, size):
138        s = 'x' * 2**26 # 64 MiB
139        f = self.MODULE.StringIO()
140        n = size
141        while n > len(s):
142            f.write(s)
143            n -= len(s)
144        s = None
145        f.write('x' * n)
146        self.assertEqual(len(f.getvalue()), size)
147
148
149class TestStringIO(TestGenericStringIO):
150    MODULE = StringIO
151
152    def test_unicode(self):
153
154        if not test_support.have_unicode: return
155
156        # The StringIO module also supports concatenating Unicode
157        # snippets to larger Unicode strings. This is tested by this
158        # method. Note that cStringIO does not support this extension.
159
160        f = self.MODULE.StringIO()
161        f.write(self._line[:6])
162        f.seek(3)
163        f.write(unicode(self._line[20:26]))
164        f.write(unicode(self._line[52]))
165        s = f.getvalue()
166        self.assertEqual(s, unicode('abcuvwxyz!'))
167        self.assertEqual(type(s), types.UnicodeType)
168
169class TestcStringIO(TestGenericStringIO):
170    MODULE = cStringIO
171
172    def test_array_support(self):
173        # Issue #1730114: cStringIO should accept array objects
174        a = array.array('B', [0,1,2])
175        f = self.MODULE.StringIO(a)
176        self.assertEqual(f.getvalue(), '\x00\x01\x02')
177
178    def test_unicode(self):
179
180        if not test_support.have_unicode: return
181
182        # The cStringIO module converts Unicode strings to character
183        # strings when writing them to cStringIO objects.
184        # Check that this works.
185
186        f = self.MODULE.StringIO()
187        f.write(u'abcde')
188        s = f.getvalue()
189        self.assertEqual(s, 'abcde')
190        self.assertEqual(type(s), str)
191
192        f = self.MODULE.StringIO(u'abcde')
193        s = f.getvalue()
194        self.assertEqual(s, 'abcde')
195        self.assertEqual(type(s), str)
196
197        self.assertRaises(UnicodeEncodeError, self.MODULE.StringIO, u'\xf4')
198
199
200import sys
201if sys.platform.startswith('java'):
202    # Jython doesn't have a buffer object, so we just do a useless
203    # fake of the buffer tests.
204    buffer = str
205
206class TestBufferStringIO(TestStringIO):
207    constructor = buffer
208
209class TestBuffercStringIO(TestcStringIO):
210    constructor = buffer
211
212class TestMemoryviewcStringIO(TestcStringIO):
213    constructor = memoryview
214
215
216def test_main():
217    test_support.run_unittest(TestStringIO, TestcStringIO)
218    with test_support.check_py3k_warnings(("buffer.. not supported",
219                                             DeprecationWarning)):
220        test_support.run_unittest(TestBufferStringIO, TestBuffercStringIO)
221    test_support.run_unittest(TestMemoryviewcStringIO)
222
223if __name__ == '__main__':
224    test_main()
225