Show More
@@ -80,16 +80,49 b' def make_line_buffered(stream):' | |||
|
80 | 80 | return LineBufferedWrapper(stream) |
|
81 | 81 | |
|
82 | 82 | |
|
83 | class WriteAllWrapper(object): | |
|
84 | def __init__(self, orig): | |
|
85 | self.orig = orig | |
|
86 | ||
|
87 | def __getattr__(self, attr): | |
|
88 | return getattr(self.orig, attr) | |
|
89 | ||
|
90 | def write(self, s): | |
|
91 | write1 = self.orig.write | |
|
92 | m = memoryview(s) | |
|
93 | total_to_write = len(s) | |
|
94 | total_written = 0 | |
|
95 | while total_written < total_to_write: | |
|
96 | total_written += write1(m[total_written:]) | |
|
97 | return total_written | |
|
98 | ||
|
99 | ||
|
100 | io.IOBase.register(WriteAllWrapper) | |
|
101 | ||
|
102 | ||
|
103 | def make_write_all(stream): | |
|
104 | assert pycompat.ispy3 | |
|
105 | if isinstance(stream, WriteAllWrapper): | |
|
106 | return stream | |
|
107 | if isinstance(stream, io.BufferedIOBase): | |
|
108 | # The io.BufferedIOBase.write() contract guarantees that all data is | |
|
109 | # written. | |
|
110 | return stream | |
|
111 | # In general, the write() method of streams is free to write only part of | |
|
112 | # the data. | |
|
113 | return WriteAllWrapper(stream) | |
|
114 | ||
|
115 | ||
|
83 | 116 | if pycompat.ispy3: |
|
84 | 117 | # Python 3 implements its own I/O streams. |
|
85 | 118 | # TODO: .buffer might not exist if std streams were replaced; we'll need |
|
86 | 119 | # a silly wrapper to make a bytes stream backed by a unicode one. |
|
87 | 120 | stdin = sys.stdin.buffer |
|
88 | stdout = sys.stdout.buffer | |
|
121 | stdout = make_write_all(sys.stdout.buffer) | |
|
89 | 122 | if isatty(stdout): |
|
90 | 123 | # The standard library doesn't offer line-buffered binary streams. |
|
91 | 124 | stdout = make_line_buffered(stdout) |
|
92 | stderr = sys.stderr.buffer | |
|
125 | stderr = make_write_all(sys.stderr.buffer) | |
|
93 | 126 | else: |
|
94 | 127 | # Python 2 uses the I/O streams provided by the C library. |
|
95 | 128 | stdin = sys.stdin |
@@ -7,6 +7,7 b' from __future__ import absolute_import' | |||
|
7 | 7 | import contextlib |
|
8 | 8 | import errno |
|
9 | 9 | import os |
|
10 | import signal | |
|
10 | 11 | import subprocess |
|
11 | 12 | import sys |
|
12 | 13 | import unittest |
@@ -31,6 +32,19 b" LINE_BUFFERED = b'[written aaa]aaabbb\\n[" | |||
|
31 | 32 | FULLY_BUFFERED = b'[written aaa][written bbb\\n]aaabbb\n' |
|
32 | 33 | |
|
33 | 34 | |
|
35 | TEST_LARGE_WRITE_CHILD_SCRIPT = r''' | |
|
36 | import signal | |
|
37 | import sys | |
|
38 | ||
|
39 | from mercurial import dispatch | |
|
40 | from mercurial.utils import procutil | |
|
41 | ||
|
42 | signal.signal(signal.SIGINT, lambda *x: None) | |
|
43 | dispatch.initstdio() | |
|
44 | procutil.{stream}.write(b'x' * 1048576) | |
|
45 | ''' | |
|
46 | ||
|
47 | ||
|
34 | 48 | @contextlib.contextmanager |
|
35 | 49 | def _closing(fds): |
|
36 | 50 | try: |
@@ -63,8 +77,8 b' def _ptys():' | |||
|
63 | 77 | yield rwpair |
|
64 | 78 | |
|
65 | 79 | |
|
66 | def _readall(fd, buffer_size): | |
|
67 | buf = [] | |
|
80 | def _readall(fd, buffer_size, initial_buf=None): | |
|
81 | buf = initial_buf or [] | |
|
68 | 82 | while True: |
|
69 | 83 | try: |
|
70 | 84 | s = os.read(fd, buffer_size) |
@@ -101,7 +115,7 b' class TestStdio(unittest.TestCase):' | |||
|
101 | 115 | ) |
|
102 | 116 | try: |
|
103 | 117 | os.close(child_stream) |
|
104 | check_output(stream_receiver) | |
|
118 | check_output(stream_receiver, proc) | |
|
105 | 119 | except: # re-raises |
|
106 | 120 | proc.terminate() |
|
107 | 121 | raise |
@@ -112,7 +126,7 b' class TestStdio(unittest.TestCase):' | |||
|
112 | 126 | def _test_buffering( |
|
113 | 127 | self, stream, rwpair_generator, expected_output, python_args=[] |
|
114 | 128 | ): |
|
115 | def check_output(stream_receiver): | |
|
129 | def check_output(stream_receiver, proc): | |
|
116 | 130 | self.assertEqual(_readall(stream_receiver, 1024), expected_output) |
|
117 | 131 | |
|
118 | 132 | self._test( |
@@ -143,6 +157,61 b' class TestStdio(unittest.TestCase):' | |||
|
143 | 157 | test_buffering_stdout_ptys_unbuffered |
|
144 | 158 | ) |
|
145 | 159 | |
|
160 | def _test_large_write(self, stream, rwpair_generator, python_args=[]): | |
|
161 | if not pycompat.ispy3 and pycompat.isdarwin: | |
|
162 | # Python 2 doesn't always retry on EINTR, but the libc might retry. | |
|
163 | # So far, it was observed only on macOS that EINTR is raised at the | |
|
164 | # Python level. As Python 2 support will be dropped soon-ish, we | |
|
165 | # won't attempt to fix it. | |
|
166 | raise unittest.SkipTest("raises EINTR on macOS") | |
|
167 | ||
|
168 | def check_output(stream_receiver, proc): | |
|
169 | if not pycompat.iswindows: | |
|
170 | # On Unix, we can provoke a partial write() by interrupting it | |
|
171 | # by a signal handler as soon as a bit of data was written. | |
|
172 | # We test that write() is called until all data is written. | |
|
173 | buf = [os.read(stream_receiver, 1)] | |
|
174 | proc.send_signal(signal.SIGINT) | |
|
175 | else: | |
|
176 | # On Windows, there doesn't seem to be a way to cause partial | |
|
177 | # writes. | |
|
178 | buf = [] | |
|
179 | self.assertEqual( | |
|
180 | _readall(stream_receiver, 131072, buf), b'x' * 1048576 | |
|
181 | ) | |
|
182 | ||
|
183 | self._test( | |
|
184 | TEST_LARGE_WRITE_CHILD_SCRIPT.format(stream=stream), | |
|
185 | stream, | |
|
186 | rwpair_generator, | |
|
187 | check_output, | |
|
188 | python_args, | |
|
189 | ) | |
|
190 | ||
|
191 | def test_large_write_stdout_pipes(self): | |
|
192 | self._test_large_write('stdout', _pipes) | |
|
193 | ||
|
194 | def test_large_write_stdout_ptys(self): | |
|
195 | self._test_large_write('stdout', _ptys) | |
|
196 | ||
|
197 | def test_large_write_stdout_pipes_unbuffered(self): | |
|
198 | self._test_large_write('stdout', _pipes, python_args=['-u']) | |
|
199 | ||
|
200 | def test_large_write_stdout_ptys_unbuffered(self): | |
|
201 | self._test_large_write('stdout', _ptys, python_args=['-u']) | |
|
202 | ||
|
203 | def test_large_write_stderr_pipes(self): | |
|
204 | self._test_large_write('stderr', _pipes) | |
|
205 | ||
|
206 | def test_large_write_stderr_ptys(self): | |
|
207 | self._test_large_write('stderr', _ptys) | |
|
208 | ||
|
209 | def test_large_write_stderr_pipes_unbuffered(self): | |
|
210 | self._test_large_write('stderr', _pipes, python_args=['-u']) | |
|
211 | ||
|
212 | def test_large_write_stderr_ptys_unbuffered(self): | |
|
213 | self._test_large_write('stderr', _ptys, python_args=['-u']) | |
|
214 | ||
|
146 | 215 | |
|
147 | 216 | if __name__ == '__main__': |
|
148 | 217 | import silenttestrunner |
General Comments 0
You need to be logged in to leave comments.
Login now