|
10 | 10 | import threading
|
11 | 11 | import uuid
|
12 | 12 | from types import TracebackType
|
13 |
| -from typing import Any, BinaryIO, Callable, Dict, List, Sequence, TextIO, Union |
| 13 | +from typing import Any, Callable, Dict, List, Sequence, TextIO, Union |
14 | 14 |
|
15 | 15 | import pydantic
|
16 | 16 | from typing_extensions import Self
|
|
19 | 19 | from .errors import CogRuntimeError, CogTimeoutError
|
20 | 20 |
|
21 | 21 |
|
22 |
| -class _SimpleStreamWrapper(io.TextIOWrapper): |
23 |
| - """ |
24 |
| - _SimpleStreamWrapper wraps a binary I/O buffer and provides a TextIOWrapper |
25 |
| - interface (primarily write and flush methods) which call a provided |
26 |
| - callback function instead of (or, if `tee` is True, in addition to) writing |
27 |
| - to the underlying buffer. |
28 |
| - """ |
29 |
| - |
30 |
| - def __init__( |
31 |
| - self, |
32 |
| - buffer: BinaryIO, |
33 |
| - callback: Callable[[str, str], None], |
34 |
| - tee: bool = False, |
35 |
| - ) -> None: |
36 |
| - super().__init__(buffer, line_buffering=True) |
37 |
| - |
38 |
| - self._callback = callback |
39 |
| - self._tee = tee |
40 |
| - self._buffer = [] |
41 |
| - |
42 |
| - def write(self, s: str) -> int: |
43 |
| - length = len(s) |
44 |
| - self._buffer.append(s) |
45 |
| - if self._tee: |
46 |
| - super().write(s) |
47 |
| - else: |
48 |
| - # If we're not teeing, we have to handle automatic flush on |
49 |
| - # newline. When `tee` is true, this is handled by the write method. |
50 |
| - if "\n" in s or "\r" in s: |
51 |
| - self.flush() |
52 |
| - return length |
53 |
| - |
54 |
| - def flush(self) -> None: |
55 |
| - self._callback(self.name, "".join(self._buffer)) |
56 |
| - self._buffer.clear() |
57 |
| - if self._tee: |
58 |
| - super().flush() |
59 |
| - |
60 |
| - |
61 | 22 | class _StreamWrapper:
|
62 | 23 | def __init__(self, name: str, stream: TextIO) -> None:
|
63 | 24 | self.name = name
|
@@ -125,66 +86,6 @@ def original(self) -> TextIO:
|
125 | 86 | return self._original_fp
|
126 | 87 |
|
127 | 88 |
|
128 |
| -if sys.version_info < (3, 9): |
129 |
| - |
130 |
| - class _AsyncStreamRedirectorBase(contextlib.AbstractContextManager): |
131 |
| - pass |
132 |
| -else: |
133 |
| - |
134 |
| - class _AsyncStreamRedirectorBase( |
135 |
| - contextlib.AbstractContextManager["AsyncStreamRedirector"] |
136 |
| - ): |
137 |
| - pass |
138 |
| - |
139 |
| - |
140 |
| -class AsyncStreamRedirector(_AsyncStreamRedirectorBase): |
141 |
| - """ |
142 |
| - AsyncStreamRedirector is a context manager that redirects I/O streams to a |
143 |
| - callback function. If `tee` is True, it also writes output to the original |
144 |
| - streams. |
145 |
| -
|
146 |
| - Unlike StreamRedirector, the underlying stream file descriptors are not |
147 |
| - modified, which means that only stream writes from Python code will be |
148 |
| - captured. Writes from native code will not be captured. |
149 |
| -
|
150 |
| - Unlike StreamRedirector, the streams redirected cannot be configured. The |
151 |
| - context manager is only able to redirect STDOUT and STDERR. |
152 |
| - """ |
153 |
| - |
154 |
| - def __init__( |
155 |
| - self, |
156 |
| - callback: Callable[[str, str], None], |
157 |
| - tee: bool = False, |
158 |
| - ) -> None: |
159 |
| - self._callback = callback |
160 |
| - self._tee = tee |
161 |
| - |
162 |
| - stdout_wrapper = _SimpleStreamWrapper(sys.stdout.buffer, callback, tee) |
163 |
| - stderr_wrapper = _SimpleStreamWrapper(sys.stderr.buffer, callback, tee) |
164 |
| - self._stdout_ctx = contextlib.redirect_stdout(stdout_wrapper) |
165 |
| - self._stderr_ctx = contextlib.redirect_stderr(stderr_wrapper) |
166 |
| - |
167 |
| - def __enter__(self) -> Self: |
168 |
| - self._stdout_ctx.__enter__() |
169 |
| - self._stderr_ctx.__enter__() |
170 |
| - return self |
171 |
| - |
172 |
| - def __exit__( |
173 |
| - self, |
174 |
| - exc_type: type[BaseException] | None, |
175 |
| - exc_value: BaseException | None, |
176 |
| - traceback: TracebackType | None, |
177 |
| - ) -> None: |
178 |
| - self._stdout_ctx.__exit__(exc_type, exc_value, traceback) |
179 |
| - self._stderr_ctx.__exit__(exc_type, exc_value, traceback) |
180 |
| - |
181 |
| - def drain(self, timeout: float = 0.0) -> None: |
182 |
| - # Draining isn't complicated for AsyncStreamRedirector, since we're not |
183 |
| - # moving data between threads. We just need to flush the streams. |
184 |
| - sys.stdout.flush() |
185 |
| - sys.stderr.flush() |
186 |
| - |
187 |
| - |
188 | 89 | if sys.version_info < (3, 9):
|
189 | 90 |
|
190 | 91 | class _StreamRedirectorBase(contextlib.AbstractContextManager):
|
|
0 commit comments