termtastic 0.0.3__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,11 @@
1
+ Metadata-Version: 2.3
2
+ Name: termtastic
3
+ Version: 0.0.3
4
+ Summary: A fantastic set of terminal utilities for python
5
+ Author: Bearmine
6
+ Requires-Python: >=3.12
7
+ Project-URL: Repository, https://forge.bearmine.com/bearmine/termtastic
8
+ Description-Content-Type: text/markdown
9
+
10
+ # termtastic
11
+
@@ -0,0 +1,2 @@
1
+ # termtastic
2
+
@@ -0,0 +1,52 @@
1
+ [project]
2
+ name = "termtastic"
3
+ version = "0.0.3"
4
+ description = "A fantastic set of terminal utilities for python"
5
+ readme = "README.md"
6
+ requires-python = ">=3.12"
7
+ dependencies = []
8
+
9
+ [[project.authors]]
10
+ name = "Bearmine"
11
+
12
+ [project.urls]
13
+ Repository = "https://forge.bearmine.com/bearmine/termtastic"
14
+
15
+ [dependency-groups]
16
+ lint = ["ruff"]
17
+ test = ["pytest"]
18
+
19
+ [[dependency-groups.dev]]
20
+ include-group = "lint"
21
+
22
+ [[dependency-groups.dev]]
23
+ include-group = "test"
24
+
25
+ [build-system]
26
+ requires = ["uv_build>=0.12.9,<0.13.0"]
27
+ build-backend = "uv_build"
28
+
29
+ [tool.ruff.lint]
30
+ extend-select = [
31
+ "F",
32
+ "W",
33
+ "E",
34
+ "I",
35
+ "UP",
36
+ "C4",
37
+ "FA",
38
+ "ISC",
39
+ "ICN",
40
+ "RET",
41
+ "SIM",
42
+ "TID",
43
+ "TC",
44
+ "PTH",
45
+ "TD",
46
+ "NPY",
47
+ "FURB",
48
+ ]
49
+ ignore = ["TD003"]
50
+
51
+ [tool.ruff.lint.pycodestyle]
52
+ max-line-length = 100
@@ -0,0 +1,59 @@
1
+ [project]
2
+ name = "termtastic"
3
+ version = "0.0.3"
4
+ description = "A fantastic set of terminal utilities for python"
5
+ readme = "README.md"
6
+ authors = [
7
+ { name = "Bearmine" }
8
+ ]
9
+ requires-python = ">=3.12"
10
+ dependencies = []
11
+
12
+ [dependency-groups]
13
+ dev = [
14
+ {include-group = "lint"},
15
+ {include-group = "test"},
16
+ ]
17
+
18
+ lint = [
19
+ "ruff"
20
+ ]
21
+
22
+ test = [
23
+ "pytest"
24
+ ]
25
+
26
+ [project.urls]
27
+ Repository = "https://forge.bearmine.com/bearmine/termtastic"
28
+
29
+ [build-system]
30
+ requires = ["uv_build>=0.12.9,<0.13.0"]
31
+ build-backend = "uv_build"
32
+
33
+ [tool.ruff.lint]
34
+ extend-select = [
35
+ "F", # Pyflakes rules
36
+ "W", # PyCodeStyle warnings
37
+ "E", # PyCodeStyle errors
38
+ "I", # Sort imports properly
39
+ "UP", # Warn if certain things can changed due to newer Python versions
40
+ "C4", # Catch incorrect use of comprehensions, dict, list, etc
41
+ "FA", # Enforce from __future__ import annotations
42
+ "ISC", # Good use of string concatenation
43
+ "ICN", # Use common import conventions
44
+ "RET", # Good return practices
45
+ "SIM", # Common simplification rules
46
+ "TID", # Some good import practices
47
+ "TC", # Enforce importing certain types in a TYPE_CHECKING block
48
+ "PTH", # Use pathlib instead of os.path
49
+ "TD", # Be diligent with TODO comments
50
+ "NPY", # Some numpy-specific things
51
+ "FURB", # Suggest more idiomatic Python patterns
52
+ ]
53
+
54
+ ignore = [
55
+ "TD003", # Missing issue link TODOs
56
+ ]
57
+
58
+ [tool.ruff.lint.pycodestyle]
59
+ max-line-length = 100
@@ -0,0 +1 @@
1
+ """A fantastic set of terminal utilities for python"""
@@ -0,0 +1,167 @@
1
+ import io
2
+ from abc import ABC, abstractmethod
3
+ from contextlib import AbstractContextManager
4
+ from types import TracebackType
5
+ from typing import Self
6
+
7
+ from termtastic import cursor
8
+ from termtastic.protected_line_writer import ProtectedLineWriter
9
+
10
+
11
+ class _AbstractBlockPrinter(ABC):
12
+ @abstractmethod
13
+ def print(self, *values, sep=" ", end="\n", flush=False):
14
+ """Print text into the block. Use this instead of normal print."""
15
+ raise NotImplementedError
16
+
17
+ @abstractmethod
18
+ def replace(self, *values, sep=" ", end="\n", flush=False):
19
+ """Replace the current block text with the given text"""
20
+ raise NotImplementedError
21
+
22
+ def flush(self):
23
+ pass
24
+
25
+ @abstractmethod
26
+ def clear(self, flush=False):
27
+ """Clear all text from the block"""
28
+ raise NotImplementedError
29
+
30
+
31
+ class _BufferedBlockPrinter(_AbstractBlockPrinter, AbstractContextManager):
32
+ def __init__(self, parent_printer: _AbstractBlockPrinter) -> None:
33
+ self._parent_printer = parent_printer
34
+ self._closed = False
35
+ self._buffer: io.StringIO = io.StringIO()
36
+
37
+ def __enter__(self) -> Self:
38
+ return self
39
+
40
+ def __exit__(
41
+ self,
42
+ exc_type: type[BaseException] | None,
43
+ exc_value: BaseException | None,
44
+ traceback: TracebackType | None,
45
+ ) -> bool | None:
46
+ self.close()
47
+
48
+ def _clear_buffer(self):
49
+ old_buffer = self._buffer
50
+ self._buffer = io.StringIO()
51
+ old_buffer.close()
52
+
53
+ def close(self):
54
+ if not self._closed:
55
+ self.flush()
56
+ self._buffer.close()
57
+ self._closed = True
58
+
59
+ def print(self, *values, sep=" ", end="\n", flush=False):
60
+ assert not self._closed, f"{self.__class__.__name__} is closed"
61
+ print(*values, sep=sep, end=end, file=self._buffer)
62
+ if flush:
63
+ self.flush()
64
+
65
+ def replace(self, *values, sep=" ", end="\n", flush=False):
66
+ # We don't need to worry about flickering at all like in BlockPrinter,
67
+ # so we can just clear and then call print.
68
+ self.clear()
69
+ self.print(*values, sep=sep, end=end, flush=flush)
70
+
71
+ def flush(self):
72
+ assert not self._closed, f"{self.__class__.__name__} is closed"
73
+ # We replace the block but do not clear the buffer as the user
74
+ # can still append additional text to the buffer
75
+ self._parent_printer.replace(self._buffer.getvalue(), end="", flush=True)
76
+
77
+ def clear(self, flush=False):
78
+ assert not self._closed, f"{self.__class__.__name__} is closed"
79
+ # Since we buffer and only write on flush, all we have to do is clear the buffer.
80
+ self._clear_buffer()
81
+ if flush:
82
+ self.flush()
83
+
84
+
85
+ class BlockPrinter(_AbstractBlockPrinter, AbstractContextManager):
86
+ def __init__(self, *, debug=False) -> None:
87
+ self._protected_printer = ProtectedLineWriter(debug=debug)
88
+ self._closed = False
89
+
90
+ def __enter__(self) -> Self:
91
+ return self
92
+
93
+ def __exit__(
94
+ self,
95
+ exc_type: type[BaseException] | None,
96
+ exc_value: BaseException | None,
97
+ traceback: TracebackType | None,
98
+ ) -> bool | None:
99
+ self.close()
100
+
101
+ def close(self):
102
+ if self._closed:
103
+ return
104
+ self._protected_printer.close()
105
+ self._closed = True
106
+
107
+ def print(self, *values, sep=" ", end="\n", flush=False):
108
+ assert not self._closed, f"{self.__class__.__name__} is closed"
109
+ print(*values, sep=sep, end=end, flush=flush, file=self._protected_printer)
110
+
111
+ def _reset_print(self, *values, sep=" ", end="\n", flush=False):
112
+ assert not self._closed, f"{self.__class__.__name__} is closed"
113
+ with self._protected_printer.lock():
114
+ self._protected_printer.reset()
115
+ print(*values, sep=sep, end=end, flush=flush, file=self._protected_printer)
116
+
117
+ def replace(self, *values, sep=" ", end="\n", flush=False):
118
+ assert not self._closed, f"{self.__class__.__name__} is closed"
119
+ with io.StringIO() as buffer:
120
+ # Return cursor to top of block
121
+ cursor.prev_line(self._protected_printer.protected_lines, file=buffer)
122
+
123
+ with io.StringIO() as str_io:
124
+ # Print given text into buffer
125
+ print(*values, sep=sep, end=end, file=str_io)
126
+
127
+ # Loop through lines clear then replace
128
+ str_io.seek(0)
129
+ for line in str_io.readlines():
130
+ if line.endswith("\n"):
131
+ print(line.rstrip("\r\n"), end="", file=buffer)
132
+ cursor.erase_to_end_of_line(file=buffer)
133
+ print(file=buffer)
134
+ else:
135
+ print(line, end="", file=buffer)
136
+ cursor.erase_to_end_of_line(file=buffer)
137
+
138
+ # Clear rest of block
139
+ cursor.clear_screen_to_end(file=buffer)
140
+
141
+ # Reset the protected printer and print buffer
142
+ self._reset_print(buffer.getvalue(), end="", flush=flush)
143
+
144
+ def clear(self, flush=False):
145
+ assert not self._closed, f"{self.__class__.__name__} is closed"
146
+ with io.StringIO() as buffer:
147
+ # Return cursor to top of block
148
+ cursor.prev_line(self._protected_printer.protected_lines, file=buffer)
149
+ # Clear all text
150
+ cursor.clear_screen_to_end(file=buffer)
151
+
152
+ # Reset the protected printer and print buffer
153
+ self._reset_print(buffer.getvalue(), end="", flush=flush)
154
+
155
+ def replace_block(self) -> _BufferedBlockPrinter:
156
+ """
157
+ Returns a buffer that can be used to replace the block.
158
+ The replacement is printed all at once on flush or close.
159
+
160
+ This can reduce text flickering when replacing using multiple prints.
161
+
162
+ Returned buffer must be closed.
163
+
164
+ Using BlockPrinter while a buffer is open is undefined behavior.
165
+ """
166
+ assert not self._closed, f"{self.__class__.__name__} is closed"
167
+ return _BufferedBlockPrinter(self)
@@ -0,0 +1,70 @@
1
+ from __future__ import annotations
2
+
3
+
4
+ class CtrlSeq:
5
+ def __init__(self, seq: str, param: None | str | int = None) -> None:
6
+ """Represents a terminal control sequence."""
7
+ self._seq = seq
8
+ self._param = "" if param is None else str(param)
9
+
10
+ def __call__(self, param: None | str | int) -> CtrlSeq:
11
+ """Set control sequence parameter"""
12
+ return CtrlSeq(self._seq, param)
13
+
14
+ def __str__(self) -> str:
15
+ return f"\033[{self._param}{self._seq}"
16
+
17
+ def __repr__(self) -> str:
18
+ return f"\\033[{self._param}{self._seq}"
19
+
20
+
21
+ CUU = CtrlSeq("A")
22
+ """Cursor Up"""
23
+
24
+ CUD = CtrlSeq("B")
25
+ """Cursor Down"""
26
+
27
+ CUF = CtrlSeq("C")
28
+ """Cursor Forward"""
29
+
30
+ CUB = CtrlSeq("D")
31
+ """Cursor Back"""
32
+
33
+ CNL = CtrlSeq("E")
34
+ """Cursor Next Line"""
35
+
36
+ CPL = CtrlSeq("F")
37
+ """Cursor Previous Line"""
38
+
39
+ CHA = CtrlSeq("G")
40
+ """Cursor Horizontal Absolute"""
41
+
42
+ CUP = CtrlSeq("H")
43
+ """Cursor Position"""
44
+
45
+ ED = CtrlSeq("J")
46
+ """Erase in Display"""
47
+
48
+ EL = CtrlSeq("K")
49
+ """Erase in Line"""
50
+
51
+ SU = CtrlSeq("S")
52
+ """Scroll Up"""
53
+
54
+ SD = CtrlSeq("T")
55
+ """Scroll Down"""
56
+
57
+ HVP = CtrlSeq("f")
58
+ """Horizontal Vertical Position"""
59
+
60
+ IL = CtrlSeq("L")
61
+ """Insert Line"""
62
+
63
+ DL = CtrlSeq("M")
64
+ """Delete Lines"""
65
+
66
+ DISPLAY_CURSOR = CtrlSeq("h")("?25")
67
+ """Make the cursor visible"""
68
+
69
+ HIDE_CURSOR = CtrlSeq("l")("?25")
70
+ """Make the cursor invisible"""
@@ -0,0 +1,207 @@
1
+ from contextlib import contextmanager
2
+
3
+ from termtastic import ctrlseq
4
+
5
+
6
+ def _clean_print(*values: object, sep=" ", file=None, flush=False):
7
+ print(*values, sep=sep, end="", file=file, flush=flush)
8
+
9
+
10
+ def flush(file=None):
11
+ _clean_print("", file=file, flush=True)
12
+
13
+
14
+ def forward(n: int, *, file=None):
15
+ """
16
+ Move cursor forwards the given number of columns
17
+
18
+ `n` must be positive
19
+ """
20
+ assert n >= 0, "n must be positive"
21
+ if n > 0:
22
+ _clean_print(ctrlseq.CUF(n), file=file)
23
+
24
+
25
+ def right(n: int, *, file=None):
26
+ """Alias for forward"""
27
+ forward(n, file=file)
28
+
29
+
30
+ def back(n: int, *, file=None):
31
+ """
32
+ Move cursor back the given number of columns
33
+
34
+ `n` must be positive
35
+ """
36
+ assert n >= 0, "n must be positive"
37
+ if n > 0:
38
+ _clean_print(ctrlseq.CUB(n), file=file)
39
+
40
+
41
+ def left(n: int, *, file=None):
42
+ """Alias for back"""
43
+ back(n, file=file)
44
+
45
+
46
+ def up(n: int, *, file=None):
47
+ """
48
+ Move cursor up the given number of lines
49
+
50
+ `n` must be positive
51
+ """
52
+ assert n >= 0, "n must be positive"
53
+ if n > 0:
54
+ _clean_print(ctrlseq.CUU(n), file=file)
55
+
56
+
57
+ def down(n: int, *, file=None):
58
+ """
59
+ Move cursor down the given number of lines
60
+
61
+ `n` must be positive
62
+ """
63
+ assert n >= 0, "n must be positive"
64
+ if n > 0:
65
+ _clean_print(ctrlseq.CUD(n), file=file)
66
+
67
+
68
+ def next_line(n: int, *, file=None):
69
+ """Move cursor to beginning of line n lines down"""
70
+ assert n >= 0, "n must be positive"
71
+ if n > 0:
72
+ _clean_print(ctrlseq.CNL(n), file=file)
73
+
74
+
75
+ def prev_line(n: int, *, file=None):
76
+ """Move cursor to beginning of line n lines up"""
77
+ assert n >= 0, "n must be positive"
78
+ if n > 0:
79
+ _clean_print(ctrlseq.CPL(n), file=file)
80
+
81
+
82
+ def horizontal_absolute(column: int, *, file=None):
83
+ """1 based column position in the current row"""
84
+ _clean_print(ctrlseq.CHA(column), file=file)
85
+
86
+
87
+ def position(row: int, column: int, *, file=None):
88
+ """Move cursor to the specified row and column (1 based)"""
89
+ _clean_print(ctrlseq.CUP(f"{row};{column}"), file=file)
90
+
91
+
92
+ def erase_display(n: int | None = None, *, file=None):
93
+ """
94
+ Clears screen.
95
+
96
+ Accepted values are:
97
+
98
+ 0 - Clear from cursor to end of screen (Default)</br>
99
+ 1 - Clear from cursor to start of screen</br>
100
+ 2 - Clear entire screen</br>
101
+ 3 - Clear entire screen and scrollback buffer
102
+ """
103
+ _clean_print(ctrlseq.ED(n), file=file)
104
+
105
+
106
+ def clear_screen_to_end(*, file=None):
107
+ erase_display(0, file=file)
108
+
109
+
110
+ def clear_screen_to_start(*, file=None):
111
+ erase_display(1, file=file)
112
+
113
+
114
+ def clear_screen(*, file=None):
115
+ erase_display(2, file=file)
116
+
117
+
118
+ def clear_screen_and_scrollback(*, file=None):
119
+ erase_display(3, file=file)
120
+
121
+
122
+ def erase_in_line(n: int | None = None, *, file=None):
123
+ """
124
+ Erases text in line. Cursor position does not change.
125
+
126
+ Accepted values are:
127
+
128
+ 0 - Clear from cursor to end of line (Default)</br>
129
+ 1 - Clear from cursor to beginning of line</br>
130
+ """
131
+ _clean_print(ctrlseq.EL(n), file=file)
132
+
133
+
134
+ def erase_to_end_of_line(*, file=None):
135
+ erase_in_line(0, file=file)
136
+
137
+
138
+ def erase_to_start_of_line(*, file=None):
139
+ erase_in_line(1, file=file)
140
+
141
+
142
+ def scroll_up(n: int, *, file=None):
143
+ """Scroll whole page up by n lines. New lines are added at the bottom."""
144
+ assert n >= 0, "n must be positive"
145
+ if n > 0:
146
+ _clean_print(ctrlseq.SU(n), file=file)
147
+
148
+
149
+ def scroll_down(n: int, *, file=None):
150
+ """Scroll whole page down by n lines. New lines are added at the top."""
151
+ assert n >= 0, "n must be positive"
152
+ if n > 0:
153
+ _clean_print(ctrlseq.SD(n), file=file)
154
+
155
+
156
+ def horizontal_vertical_position(row: int, column: int, *, file=None):
157
+ _clean_print(ctrlseq.HVP(f"{row};{column}"))
158
+
159
+
160
+ def insert_lines(n: int, *, file=None):
161
+ """
162
+ Insert `n` lines at current cursor position.
163
+
164
+ `n` must be positive
165
+ """
166
+ assert n >= 0, "n must be positive"
167
+ if n > 0:
168
+ _clean_print(ctrlseq.IL(n), file=file)
169
+
170
+
171
+ def delete_lines(n: int, *, file=None):
172
+ """
173
+ Delete `n` lines at current cursor position.
174
+
175
+ `n` must be positive
176
+ """
177
+ assert n >= 0, "n must be positive"
178
+ if n > 0:
179
+ _clean_print(ctrlseq.DL(n), file=file)
180
+
181
+
182
+ def enable_cursor(*, file=None):
183
+ """Enable display of cursor"""
184
+ _clean_print(ctrlseq.DISPLAY_CURSOR, file=file, flush=True)
185
+
186
+
187
+ def disable_cursor(*, file=None):
188
+ """Disable display of cursor"""
189
+ _clean_print(ctrlseq.HIDE_CURSOR, file=file, flush=True)
190
+
191
+
192
+ @contextmanager
193
+ def show_cursor(*, file=None):
194
+ enable_cursor(file=file)
195
+ try:
196
+ yield
197
+ finally:
198
+ disable_cursor(file=file)
199
+
200
+
201
+ @contextmanager
202
+ def hide_cursor(*, file=None):
203
+ disable_cursor(file=file)
204
+ try:
205
+ yield
206
+ finally:
207
+ enable_cursor(file=file)
@@ -0,0 +1,350 @@
1
+ import logging
2
+ import math
3
+ import shutil
4
+ import sys
5
+ import time
6
+ from abc import abstractmethod
7
+ from contextlib import AbstractContextManager, contextmanager
8
+ from io import StringIO, TextIOBase
9
+ from threading import RLock
10
+ from types import TracebackType
11
+ from typing import Literal, Self, TextIO
12
+
13
+ from termtastic import cursor
14
+
15
+
16
+ def _get_all_stream_handlers() -> set[logging.StreamHandler]:
17
+ handler_set: set[logging.StreamHandler] = set()
18
+ for handler in logging.root.handlers:
19
+ if isinstance(handler, logging.StreamHandler):
20
+ handler_set.add(handler)
21
+
22
+ children = logging.root.getChildren()
23
+
24
+ while children:
25
+ sub_children = set()
26
+ for child in children:
27
+ for handler in child.handlers:
28
+ if isinstance(handler, logging.StreamHandler):
29
+ handler_set.add(handler)
30
+ sub_children.update(child.getChildren())
31
+ children = sub_children
32
+
33
+ return handler_set
34
+
35
+
36
+ def _newlines_in_terminal(text: str, columns: int) -> int:
37
+ """
38
+ Return how many lines in a terminal with the given number of columns the given text will take up
39
+ """
40
+ if len(text) == 0:
41
+ return 0
42
+ newlines = text.count("\n")
43
+ for ln in text.splitlines():
44
+ newlines += math.floor(max(0, len(ln) - 1) / (columns))
45
+ return newlines
46
+
47
+
48
+ class _StreamRedirect[T](AbstractContextManager):
49
+ def __init__(self, old_stream: T) -> None:
50
+ super().__init__()
51
+ self._old_stream: T = old_stream
52
+ self._closed = False
53
+
54
+ @property
55
+ def old_stream(self) -> T:
56
+ return self._old_stream
57
+
58
+ def close(self):
59
+ if self._closed:
60
+ return
61
+ self._close()
62
+ self._closed = True
63
+
64
+ @abstractmethod
65
+ def _close(self):
66
+ raise NotImplementedError
67
+
68
+ def __enter__(self) -> Self:
69
+ return self
70
+
71
+ def __exit__(
72
+ self,
73
+ exc_type: type[BaseException] | None,
74
+ exc_value: BaseException | None,
75
+ traceback: TracebackType | None,
76
+ ) -> bool | None:
77
+ self.close()
78
+
79
+
80
+ class _RedirectStreamHandler(_StreamRedirect):
81
+ def __init__(self, handler: logging.StreamHandler, stream: TextIOBase) -> None:
82
+ super().__init__(handler.setStream(stream))
83
+ self._handler = handler
84
+
85
+ def _close(self):
86
+ if self.old_stream:
87
+ self._handler.setStream(self.old_stream)
88
+
89
+
90
+ class _RedirectStdout(_StreamRedirect[TextIO]):
91
+ def __init__(self, stream: TextIOBase) -> None:
92
+ """Redirect sys.stdout to the given stream"""
93
+ super().__init__(sys.stdout)
94
+ sys.stdout = stream
95
+
96
+ def _close(self):
97
+ sys.stdout = self.old_stream
98
+
99
+
100
+ class _RedirectStderr(_StreamRedirect[TextIO]):
101
+ def __init__(self, stream: TextIOBase) -> None:
102
+ """Redirect sys.stderr to the given stream"""
103
+ super().__init__(sys.stderr)
104
+ sys.stderr = stream
105
+
106
+ def _close(self):
107
+ sys.stderr = self.old_stream
108
+
109
+
110
+ class ProtectedLineWriter(TextIOBase, AbstractContextManager):
111
+ def __init__(
112
+ self,
113
+ io: TextIOBase | TextIO | None = None,
114
+ *,
115
+ redirect_stdout=True,
116
+ redirect_stderr=True,
117
+ redirect_logging=True,
118
+ debug=False,
119
+ debug_delay_sec=0.25,
120
+ ) -> None:
121
+ super().__init__()
122
+ self._io = io or sys.stdout
123
+ self._last_writer: TextIOBase | TextIO | None = None
124
+ self._protected_lines: int = 0
125
+ self._prev_line_len: int = 0
126
+ self._buffered_newline: bool = True
127
+ self._debug_delay_sec = debug_delay_sec
128
+ self._closed = False
129
+
130
+ # Print lock used for protecting against
131
+ # multithreaded print statements
132
+ self._print_lock = RLock()
133
+
134
+ # Setup redirects
135
+ writer_cls = (
136
+ _ProtectedRedirectWriterDEBUG if debug else _ProtectedRedirectWriter
137
+ )
138
+
139
+ self._redirects: list[_StreamRedirect[TextIO]] = []
140
+
141
+ if redirect_stdout:
142
+ self._stdout_writer = writer_cls(self, sys.stdout)
143
+ self._redirects.extend(
144
+ self._redirect_sys_stream(
145
+ "stdout", self._stdout_writer, redirect_logging
146
+ )
147
+ )
148
+ self._stdout_writer.flush()
149
+ else:
150
+ self._stdout_writer = None
151
+
152
+ if redirect_stderr:
153
+ self._stderr_writer = writer_cls(self, sys.stderr)
154
+ self._redirects.extend(
155
+ self._redirect_sys_stream(
156
+ "stderr", self._stderr_writer, redirect_logging
157
+ )
158
+ )
159
+ self._stderr_writer.flush()
160
+ else:
161
+ self._stderr_writer = None
162
+
163
+ def _redirect_sys_stream(
164
+ self,
165
+ sys_stream: Literal["stdout", "stderr"],
166
+ writer: TextIOBase,
167
+ redirect_logging: bool,
168
+ ) -> list[_StreamRedirect[TextIO]]:
169
+ if sys_stream == "stdout":
170
+ sys_redirect = _RedirectStdout(writer)
171
+ else:
172
+ sys_redirect = _RedirectStderr(writer)
173
+
174
+ stream_redirects: list[_StreamRedirect[TextIO]] = []
175
+ stream_redirects.append(sys_redirect)
176
+
177
+ if redirect_logging:
178
+ for handler in _get_all_stream_handlers():
179
+ if handler.stream is sys_redirect.old_stream:
180
+ stream_redirects.append(_RedirectStreamHandler(handler, writer))
181
+
182
+ return stream_redirects
183
+
184
+ def __enter__(self) -> Self:
185
+ return self
186
+
187
+ def __exit__(
188
+ self,
189
+ exc_type: type[BaseException] | None,
190
+ exc_val: BaseException | None,
191
+ exc_tb: TracebackType | None,
192
+ ) -> None:
193
+ self.close()
194
+
195
+ @contextmanager
196
+ def lock(self, writer: TextIO | TextIOBase | None = None):
197
+ """Acquire lock and flush previous writer if necessary"""
198
+ try:
199
+ self._print_lock.acquire()
200
+ if writer and self._last_writer and self._last_writer is not writer:
201
+ self._last_writer.flush()
202
+ yield
203
+ finally:
204
+ if writer:
205
+ self._last_writer = writer
206
+ self._print_lock.release()
207
+
208
+ @property
209
+ def protected_lines(self):
210
+ return self._protected_lines
211
+
212
+ def close(self):
213
+ if self._closed:
214
+ return
215
+
216
+ self.flush()
217
+
218
+ # Close all redirects
219
+ for redirect in self._redirects:
220
+ redirect.close()
221
+
222
+ # Close stdout and stderr writers
223
+ if self._stdout_writer:
224
+ self._stdout_writer.close()
225
+
226
+ if self._stderr_writer:
227
+ self._stderr_writer.close()
228
+
229
+ self._closed = True
230
+
231
+ def flush(self):
232
+ with self.lock(self._io):
233
+ self._io.flush()
234
+
235
+ def write(self, s: str) -> int:
236
+ assert not self._closed, f"{ProtectedLineWriter.__class__.__name__} is closed"
237
+ with self.lock(self._io):
238
+ # Add number of line breaks to protected lines
239
+ self._protected_lines += _newlines_in_terminal(
240
+ s, shutil.get_terminal_size().columns
241
+ )
242
+ self._io.write(s)
243
+ return len(s)
244
+
245
+ def reset(self):
246
+ with self.lock(self._io):
247
+ self._protected_lines = 0
248
+ self._prev_line_len = 0
249
+
250
+
251
+ class _ProtectedRedirectWriter(TextIOBase, AbstractContextManager):
252
+ def __init__(
253
+ self,
254
+ printer: ProtectedLineWriter,
255
+ io: TextIOBase | TextIO,
256
+ ) -> None:
257
+ super().__init__()
258
+ self._printer = printer
259
+ self._io = io
260
+
261
+ def _debug_delay(self, buffer: StringIO):
262
+ return
263
+
264
+ def flush(self) -> None:
265
+ with self._printer.lock(self._io):
266
+ self._io.flush()
267
+
268
+ def close(self) -> None:
269
+ self.flush()
270
+
271
+ def _insert_lines(self, lines_to_insert: int, lines_up: int, buffer: StringIO):
272
+ # TODO(Bearmine): Has issues at edges of screen. i.e. if lines to insert is greats that
273
+ # number of lines in the terminal
274
+
275
+ # Move all text up to make room for new lines
276
+ # If we don't do this we lose the text at the bottom of the terminal
277
+ print("\n" * lines_to_insert, end="", file=buffer)
278
+ # Move cursor up to position to insert
279
+ cursor.prev_line(lines_up + lines_to_insert, file=buffer)
280
+ # Insert blank lines (this will push existing text down)
281
+ cursor.insert_lines(lines_to_insert, file=buffer)
282
+
283
+ def write(self, s: str) -> int:
284
+ # Empty String
285
+ if not s:
286
+ return 0
287
+
288
+ # Write out string while protecting lines
289
+ with self._printer.lock(self._io):
290
+ with StringIO() as buffer:
291
+ term_size = shutil.get_terminal_size()
292
+ lines_to_insert = _newlines_in_terminal(
293
+ s.rstrip("\r\n"), term_size.columns
294
+ )
295
+
296
+ if self._printer._buffered_newline:
297
+ lines_to_insert += 1
298
+
299
+ self._insert_lines(
300
+ lines_to_insert, self._printer.protected_lines, buffer=buffer
301
+ )
302
+ self._debug_delay(buffer)
303
+
304
+ # Move cursor up one because we should be 1 more line up
305
+ # before we print, as we are currently on a blank inserted line.
306
+ #
307
+ # We might actually be on a partial line.
308
+ cursor.prev_line(1, file=buffer)
309
+
310
+ # Move the cursor into position (end of last write)
311
+ cursor.forward(self._printer._prev_line_len, file=buffer)
312
+ self._debug_delay(buffer)
313
+
314
+ # Fill blank space with given text
315
+ if self._printer._buffered_newline:
316
+ buffer.write("\n")
317
+ self._printer._buffered_newline = False
318
+ buffer.write(s.rstrip("\r\n"))
319
+ self._debug_delay(buffer)
320
+
321
+ # Move back to bottom of terminal
322
+ cursor.next_line(self._printer._protected_lines + 1, file=buffer)
323
+ self._debug_delay(buffer)
324
+
325
+ # Track lines that don't end in newline so we can place the cursor correctly the
326
+ # next time a write occurs
327
+ if s.endswith("\n"):
328
+ self._printer._prev_line_len = 0
329
+ self._printer._buffered_newline = True
330
+ else:
331
+ if lines_to_insert == 0:
332
+ self._printer._prev_line_len += len(s) % term_size.columns
333
+ else:
334
+ self._printer._prev_line_len = (
335
+ len(s.splitlines()[-1]) % term_size.columns
336
+ )
337
+
338
+ # Write result
339
+ self._io.write(buffer.getvalue())
340
+
341
+ return len(s)
342
+
343
+
344
+ class _ProtectedRedirectWriterDEBUG(_ProtectedRedirectWriter):
345
+ def _debug_delay(self, buffer: StringIO):
346
+ self._io.write(buffer.getvalue())
347
+ buffer.seek(0)
348
+ buffer.truncate()
349
+ self.flush()
350
+ time.sleep(self._printer._debug_delay_sec)