-
Notifications
You must be signed in to change notification settings - Fork 12
/
streams.py
313 lines (238 loc) · 8.27 KB
/
streams.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
import io
import sys
import time
from ctypes import byref, c_ulong
from .buffer import get_buffer
from .info import WINDOWS, PY2
if PY2:
from .file_object import FileObject
if WINDOWS:
from ctypes import windll
kernel32 = windll.kernel32
GetStdHandle = kernel32.GetStdHandle
ReadConsoleW = kernel32.ReadConsoleW
WriteConsoleW = kernel32.WriteConsoleW
GetLastError = kernel32.GetLastError
STDIN_HANDLE = GetStdHandle(-10)
STDOUT_HANDLE = GetStdHandle(-11)
STDERR_HANDLE = GetStdHandle(-12)
ERROR_SUCCESS = 0
ERROR_NOT_ENOUGH_MEMORY = 8
ERROR_OPERATION_ABORTED = 995
STDIN_FILENO = 0
STDOUT_FILENO = 1
STDERR_FILENO = 2
EOF = b"\x1a"
MAX_BYTES_WRITTEN = 32767 # arbitrary because WriteConsoleW ability to write big buffers depends on heap usage
class _ReprMixin:
def __repr__(self):
modname = self.__class__.__module__
if PY2:
clsname = self.__class__.__name__
else:
clsname = self.__class__.__qualname__
attributes = []
for name in ["name", "encoding"]:
try:
value = getattr(self, name)
except AttributeError:
pass
else:
attributes.append("{}={}".format(name, repr(value)))
return "<{}.{} {}>".format(modname, clsname, " ".join(attributes))
class WindowsConsoleRawIOBase(_ReprMixin, io.RawIOBase):
def __init__(self, name, handle, fileno):
self.name = name
self.handle = handle
self.file_no = fileno
def fileno(self):
return self.file_no
def isatty(self):
# PY3 # super().isatty() # for close check in default implementation
super(WindowsConsoleRawIOBase, self).isatty()
return True
class WindowsConsoleRawReader(WindowsConsoleRawIOBase):
def readable(self):
return True
def readinto(self, b):
bytes_to_be_read = len(b)
if not bytes_to_be_read:
return 0
elif bytes_to_be_read % 2:
raise ValueError("cannot read odd number of bytes from UTF-16-LE encoded console")
buffer = get_buffer(b, writable=True)
code_units_to_be_read = bytes_to_be_read // 2
code_units_read = c_ulong()
retval = ReadConsoleW(self.handle, buffer, code_units_to_be_read, byref(code_units_read), None)
if GetLastError() == ERROR_OPERATION_ABORTED:
time.sleep(0.1) # wait for KeyboardInterrupt
if not retval:
raise OSError("Windows error {}".format(GetLastError()))
if buffer[0] == EOF:
return 0
else:
return 2 * code_units_read.value
class WindowsConsoleRawWriter(WindowsConsoleRawIOBase):
def writable(self):
return True
@staticmethod
def _error_message(errno):
if errno == ERROR_SUCCESS:
return "Windows error {} (ERROR_SUCCESS); zero bytes written on nonzero input, probably just one byte given".format(errno)
elif errno == ERROR_NOT_ENOUGH_MEMORY:
return "Windows error {} (ERROR_NOT_ENOUGH_MEMORY); try to lower `win_unicode_console.streams.MAX_BYTES_WRITTEN`".format(errno)
else:
return "Windows error {}".format(errno)
def write(self, b):
bytes_to_be_written = len(b)
buffer = get_buffer(b)
code_units_to_be_written = min(bytes_to_be_written, MAX_BYTES_WRITTEN) // 2
code_units_written = c_ulong()
retval = WriteConsoleW(self.handle, buffer, code_units_to_be_written, byref(code_units_written), None)
bytes_written = 2 * code_units_written.value
# fixes both infinite loop of io.BufferedWriter.flush() on when the buffer has odd length
# and situation when WriteConsoleW refuses to write lesser than MAX_BYTES_WRITTEN bytes
if bytes_written == 0 != bytes_to_be_written:
raise OSError(self._error_message(GetLastError()))
else:
return bytes_written
class _TextStreamWrapperMixin(_ReprMixin):
def __init__(self, base):
self.base = base
@property
def encoding(self):
return self.base.encoding
@property
def errors(self):
return self.base.errors
@property
def line_buffering(self):
return self.base.line_buffering
def seekable(self):
return self.base.seekable()
def readable(self):
return self.base.readable()
def writable(self):
return self.base.writable()
def flush(self):
self.base.flush()
def close(self):
self.base.close()
@property
def closed(self):
return self.base.closed
@property
def name(self):
return self.base.name
def fileno(self):
return self.base.fileno()
def isatty(self):
return self.base.isatty()
def write(self, s):
return self.base.write(s)
def tell(self):
return self.base.tell()
def truncate(self, pos=None):
return self.base.truncate(pos)
def seek(self, cookie, whence=0):
return self.base.seek(cookie, whence)
def read(self, size=None):
return self.base.read(size)
def __next__(self):
return next(self.base)
def readline(self, size=-1):
return self.base.readline(size)
@property
def newlines(self):
return self.base.newlines
class TextStreamWrapper(_TextStreamWrapperMixin, io.TextIOBase):
pass
class TextTranscodingWrapper(TextStreamWrapper):
encoding = None # disable the descriptor
def __init__(self, base, encoding):
# PY3 # super().__init__(base)
super(TextTranscodingWrapper, self).__init__(base)
self.encoding = encoding
class StrStreamWrapper(TextStreamWrapper):
def write(self, s):
if isinstance(s, bytes):
s = s.decode(self.encoding)
self.base.write(s)
if PY2:
class FileobjWrapper(_TextStreamWrapperMixin, file):
def __init__(self, base, f):
super(FileobjWrapper, self).__init__(base)
fileobj = self._fileobj = FileObject.from_file(self)
fileobj.set_encoding(base.encoding)
fileobj.copy_file_pointer(f)
fileobj.readable = base.readable()
fileobj.writable = base.writable()
# needed for the right interpretation of unicode literals in interactive mode when win_unicode_console is enabled in sitecustomize since Py_Initialize changes encoding afterwards
def _reset_encoding(self):
self._fileobj.set_encoding(self.base.encoding)
def readline(self, size=-1):
self._reset_encoding()
return self.base.readline(size)
if WINDOWS:
stdin_raw = WindowsConsoleRawReader("<stdin>", STDIN_HANDLE, STDIN_FILENO)
stdout_raw = WindowsConsoleRawWriter("<stdout>", STDOUT_HANDLE, STDOUT_FILENO)
stderr_raw = WindowsConsoleRawWriter("<stderr>", STDERR_HANDLE, STDERR_FILENO)
stdin_text = io.TextIOWrapper(io.BufferedReader(stdin_raw), encoding="utf-16-le", line_buffering=True)
stdout_text = io.TextIOWrapper(io.BufferedWriter(stdout_raw), encoding="utf-16-le", line_buffering=True)
stderr_text = io.TextIOWrapper(io.BufferedWriter(stderr_raw), encoding="utf-16-le", line_buffering=True)
stdin_text_transcoded = TextTranscodingWrapper(stdin_text, encoding="utf-8")
stdout_text_transcoded = TextTranscodingWrapper(stdout_text, encoding="utf-8")
stderr_text_transcoded = TextTranscodingWrapper(stderr_text, encoding="utf-8")
stdout_text_str = StrStreamWrapper(stdout_text_transcoded)
stderr_text_str = StrStreamWrapper(stderr_text_transcoded)
if PY2:
stdin_text_fileobj = FileobjWrapper(stdin_text_transcoded, sys.__stdin__)
stdout_text_str_fileobj = FileobjWrapper(stdout_text_str, sys.__stdout__)
def disable():
sys.stdin.flush()
sys.stdout.flush()
sys.stderr.flush()
sys.stdin = sys.__stdin__
sys.stdout = sys.__stdout__
sys.stderr = sys.__stderr__
def check_stream(stream, fileno):
if stream is None: # e.g. with IDLE
return True
try:
_fileno = stream.fileno()
except io.UnsupportedOperation:
return False
else:
if _fileno == fileno and stream.isatty():
stream.flush()
return True
else:
return False
# PY3 # def enable(*, stdin=Ellipsis, stdout=Ellipsis, stderr=Ellipsis):
def enable(stdin=Ellipsis, stdout=Ellipsis, stderr=Ellipsis):
if not WINDOWS:
return
# defaults
if PY2:
if stdin is Ellipsis:
stdin = stdin_text_fileobj
if stdout is Ellipsis:
stdout = stdout_text_str
if stderr is Ellipsis:
stderr = stderr_text_str
else: # transcoding because Python tokenizer cannot handle UTF-16
if stdin is Ellipsis:
stdin = stdin_text_transcoded
if stdout is Ellipsis:
stdout = stdout_text_transcoded
if stderr is Ellipsis:
stderr = stderr_text_transcoded
if stdin is not None and check_stream(sys.stdin, STDIN_FILENO):
sys.stdin = stdin
if stdout is not None and check_stream(sys.stdout, STDOUT_FILENO):
sys.stdout = stdout
if stderr is not None and check_stream(sys.stderr, STDERR_FILENO):
sys.stderr = stderr
# PY3 # def enable_only(*, stdin=None, stdout=None, stderr=None):
def enable_only(stdin=None, stdout=None, stderr=None):
enable(stdin=stdin, stdout=stdout, stderr=stderr)