66 lines
1.9 KiB
Python
66 lines
1.9 KiB
Python
import io
|
|
from contextlib import contextmanager
|
|
|
|
from ..util import trim
|
|
|
|
|
|
class IndentPrintMixin:
|
|
def __init__(self, default_indent: int = 4):
|
|
self.default_indent = default_indent
|
|
self.indent_stack: list[int] = [0]
|
|
self.output_stream = io.StringIO()
|
|
|
|
def printed_text(self):
|
|
return self.output_stream.getvalue()
|
|
|
|
def print(self, *args, **kwargs):
|
|
args = [trim(str(arg)) for arg in args if arg is not None]
|
|
lines = self.as_printed_lines(*args, **kwargs)
|
|
self._do_print_lines(lines, **kwargs)
|
|
|
|
def as_printed(self, *args, **kwargs):
|
|
lines = self.as_printed_lines(*args, **kwargs)
|
|
return '\n'.join(lines)
|
|
|
|
def as_printed_lines(self, *args, **kwargs):
|
|
text = self.io_print(*args, **kwargs)
|
|
return self.indented_lines(text)
|
|
|
|
@contextmanager
|
|
def indent(self, amount: int | None = None):
|
|
assert amount is None or amount >= 0
|
|
if amount is None:
|
|
amount = self.default_indent
|
|
|
|
self.indent_stack.append(amount + self.indent_stack[-1])
|
|
try:
|
|
yield
|
|
finally:
|
|
self.indent_stack.pop()
|
|
|
|
@property
|
|
def current_indentation(self):
|
|
return ' ' * self.indent_stack[-1]
|
|
|
|
@staticmethod
|
|
def io_print(*args, **kwargs):
|
|
kwargs.pop('file', None) # do not allow redirection of output
|
|
with io.StringIO() as output:
|
|
print(*args, file=output, **kwargs)
|
|
return output.getvalue()
|
|
|
|
def _do_print_lines(self, lines: list[str] | None = None, **kwargs):
|
|
if not lines:
|
|
print(file=self.output_stream, **kwargs)
|
|
return
|
|
|
|
for line in lines:
|
|
print(line, file=self.output_stream, **kwargs)
|
|
|
|
def indented_lines(self, text):
|
|
text = trim(text)
|
|
return [
|
|
(self.current_indentation + line).rstrip()
|
|
for line in text.splitlines(keepends=False)
|
|
]
|