from __future__ import annotations
import atexit
import contextlib
import datetime
import io
import logging
import os
import re
import sys
import typing
from collections.abc import Callable, Iterable, Iterator, Mapping
from types import TracebackType
from python_utils import types
from python_utils.converters import scale_1024
from python_utils.terminal import get_terminal_size
from python_utils.time import epoch, format_time, timedelta_to_seconds
from progressbar import base, env, terminal
# Make sure these are available for import
assert timedelta_to_seconds is not None
assert get_terminal_size is not None
assert format_time is not None
assert scale_1024 is not None
assert epoch is not None
StringT = typing.TypeVar('StringT', bound=types.StringTypes)
T = typing.TypeVar('T')
logger: logging.Logger = logging.getLogger(__name__)
class _ProgressListener(typing.Protocol):
"""Structural type for the bars the stream wrapper notifies.
Defined locally instead of importing ``ProgressBarMixinBase`` from
``bar`` so ``utils`` has no dependency on ``bar`` — not even a
type-checking-only one, which CodeQL still reports as a ``bar`` <->
``utils`` module-level import cycle.
"""
def update(self) -> None:
"""Redraw in response to redirected output being written."""
# Precompiled ANSI CSI escape-sequence patterns (str and bytes). Compiled once
# at import instead of per no_color() call, which runs for every widget on
# every redraw.
_ANSI_COLOR_RE: re.Pattern[str] = re.compile('\x1b\\[.*?[@-~]')
_ANSI_COLOR_RE_BYTES: re.Pattern[bytes] = re.compile(
bytes(terminal.ESC, 'ascii') + b'\\[.*?[@-~]',
)
@typing.overload
def deltas_to_seconds(
*deltas: None | datetime.timedelta | float | int,
default: type[ValueError] = ...,
) -> float:
"""Coalesce to seconds; raise ``ValueError`` if no delta is valid."""
@typing.overload
def deltas_to_seconds(
*deltas: None | datetime.timedelta | float | int,
default: T,
) -> float | T:
"""Coalesce to seconds; return ``default`` if no delta is valid."""
[docs]
def deltas_to_seconds(
*deltas: None | datetime.timedelta | float | int,
default: typing.Any = ValueError,
) -> typing.Any:
"""
Convert timedeltas and seconds as int to seconds as float while coalescing.
>>> deltas_to_seconds(datetime.timedelta(seconds=1, milliseconds=234))
1.234
>>> deltas_to_seconds(123)
123.0
>>> deltas_to_seconds(1.234)
1.234
>>> deltas_to_seconds(None, 1.234)
1.234
>>> deltas_to_seconds(0, 1.234)
0.0
>>> deltas_to_seconds()
Traceback (most recent call last):
...
ValueError: No valid deltas passed to `deltas_to_seconds`
>>> deltas_to_seconds(None)
Traceback (most recent call last):
...
ValueError: No valid deltas passed to `deltas_to_seconds`
>>> deltas_to_seconds(default=0.0)
0.0
"""
for delta in deltas:
if delta is None:
continue
if isinstance(delta, datetime.timedelta):
return timedelta_to_seconds(delta)
elif not isinstance(delta, float):
return float(delta)
else:
return delta
if default is ValueError:
raise ValueError('No valid deltas passed to `deltas_to_seconds`')
else:
return default
[docs]
def no_color(value: StringT) -> StringT:
"""
Return the `value` without ANSI escape codes.
>>> no_color(b'\u001b[1234]abc')
b'abc'
>>> str(no_color('\u001b[1234]abc'))
'abc'
>>> str(no_color('\u001b[1234]abc'))
'abc'
>>> no_color(123)
Traceback (most recent call last):
...
TypeError: `value` must be a string or bytes, got 123
"""
if isinstance(value, bytes):
# Fast path: with no ESC byte there is nothing to strip, so the regex
# would return the value unchanged anyway. Skipping it avoids a
# substitution on the common plain-text case, which dominates the
# per-redraw render cost (len_color is called for every widget).
if b'\x1b' not in value:
return value # type: ignore
return _ANSI_COLOR_RE_BYTES.sub(b'', value) # type: ignore
elif isinstance(value, str):
if '\x1b' not in value:
return value # type: ignore
return _ANSI_COLOR_RE.sub('', value) # type: ignore
else:
raise TypeError(f'`value` must be a string or bytes, got {value!r}')
[docs]
def len_color(value: types.StringTypes) -> int:
"""
Return the length of `value` without ANSI escape codes.
>>> len_color(b'\u001b[1234]abc')
3
>>> len_color('\u001b[1234]abc')
3
>>> len_color('\u001b[1234]abc')
3
"""
return len(no_color(value))
[docs]
class WrappingIO:
buffer: io.StringIO
target: base.IO
capturing: bool
listeners: set[_ProgressListener]
needs_clear: bool = False
def __init__(
self,
target: base.IO,
capturing: bool = False,
listeners: set[_ProgressListener] | None = None,
) -> None:
self.buffer = io.StringIO()
self.target = target
self.capturing = capturing
self.listeners = listeners or set()
self.needs_clear = False
[docs]
def write(self, value: str) -> int:
ret = 0
if self.capturing:
ret += self.buffer.write(value)
if '\n' in value: # pragma: no branch
self.needs_clear = True
for listener in self.listeners: # pragma: no branch
listener.update()
else:
ret += self.target.write(value)
if '\n' in value: # pragma: no branch
self.flush_target()
return ret
[docs]
def flush(self) -> None:
self.buffer.flush()
def _flush(self) -> None:
if value := self.buffer.getvalue():
self.flush()
# Clear the buffer before writing so a failed write cannot
# cause the same data to be written again by the next flush
self.buffer.seek(0)
self.buffer.truncate(0)
self.needs_clear = False
if not self.target.closed:
self.target.write(value)
# when explicitly flushing, always flush the target as well
self.flush_target()
[docs]
def flush_target(self) -> None: # pragma: no cover
if not self.target.closed and getattr(self.target, 'flush', None):
self.target.flush()
def __enter__(self) -> WrappingIO:
return self
[docs]
def fileno(self) -> int:
return self.target.fileno()
[docs]
def isatty(self) -> bool:
return self.target.isatty()
[docs]
def read(self, n: int = -1) -> str:
return self.target.read(n)
[docs]
def readable(self) -> bool:
return self.target.readable()
[docs]
def readline(self, limit: int = -1) -> str:
return self.target.readline(limit)
[docs]
def readlines(self, hint: int = -1) -> list[str]:
return self.target.readlines(hint)
[docs]
def seek(self, offset: int, whence: int = os.SEEK_SET) -> int:
return self.target.seek(offset, whence)
[docs]
def seekable(self) -> bool:
return self.target.seekable()
[docs]
def tell(self) -> int:
return self.target.tell()
[docs]
def truncate(self, size: int | None = None) -> int:
return self.target.truncate(size)
[docs]
def writable(self) -> bool:
return self.target.writable()
[docs]
def writelines(self, lines: Iterable[str]) -> None:
return self.target.writelines(lines)
[docs]
def close(self) -> None:
self.flush()
self.target.close()
def __next__(self) -> str:
return self.target.__next__()
def __iter__(self) -> Iterator[str]:
return self.target.__iter__()
def __exit__(
self,
__t: type[BaseException] | None,
__value: BaseException | None,
__traceback: TracebackType | None,
) -> None:
self.close()
[docs]
class StreamWrapper:
"""Wrap stdout and stderr globally."""
stdout: base.TextIO | WrappingIO
stderr: base.TextIO | WrappingIO
original_excepthook: Callable[
[
type[BaseException],
BaseException,
TracebackType | None,
],
None,
]
wrapped_stdout: int = 0
wrapped_stderr: int = 0
wrapped_logging: int = 0
wrapped_excepthook: int = 0
logging_handlers: list[tuple[logging.StreamHandler[base.IO], base.IO]]
capturing: int = 0
listeners: set[_ProgressListener]
def __init__(self) -> None:
self.stdout = self.original_stdout = sys.stdout
self.stderr = self.original_stderr = sys.stderr
self.original_excepthook = sys.excepthook
self.wrapped_stdout = 0
self.wrapped_stderr = 0
self.wrapped_logging = 0
self.wrapped_excepthook = 0
self.logging_handlers = []
self.capturing = 0
self.listeners = set()
if env.env_flag('WRAP_STDOUT', default=False): # pragma: no cover
self.wrap_stdout()
if env.env_flag('WRAP_STDERR', default=False): # pragma: no cover
self.wrap_stderr()
[docs]
def start_capturing(self, bar: _ProgressListener | None = None) -> None:
if bar: # pragma: no branch
self.listeners.add(bar)
self.capturing += 1
self.update_capturing()
[docs]
def stop_capturing(self, bar: _ProgressListener | None = None) -> None:
if bar: # pragma: no branch
with contextlib.suppress(KeyError):
self.listeners.remove(bar)
self.capturing -= 1
self.update_capturing()
[docs]
def update_capturing(self) -> None: # pragma: no cover
if isinstance(self.stdout, WrappingIO):
self.stdout.capturing = self.capturing > 0
if isinstance(self.stderr, WrappingIO):
self.stderr.capturing = self.capturing > 0
if self.capturing <= 0:
self.flush()
[docs]
def wrap(self, stdout: bool = False, stderr: bool = False) -> None:
if stdout:
self.wrap_stdout()
if stderr:
self.wrap_stderr()
[docs]
def wrap_stdout(self) -> WrappingIO:
self.wrap_excepthook()
if not self.wrapped_stdout:
self.stdout = sys.stdout = WrappingIO( # type: ignore
self.original_stdout,
listeners=self.listeners,
)
self.wrapped_stdout += 1
return sys.stdout # type: ignore
[docs]
def wrap_stderr(self) -> WrappingIO:
self.wrap_excepthook()
if not self.wrapped_stderr:
self.stderr = sys.stderr = WrappingIO( # type: ignore
self.original_stderr,
listeners=self.listeners,
)
self.wrapped_stderr += 1
return sys.stderr # type: ignore
[docs]
def wrap_logging(self) -> None:
"""Retarget stdout/stderr logging handlers to wrapped streams."""
self.wrapped_logging += 1
if self.wrapped_logging > 1:
return
wrapped_streams = {
self.original_stdout: self.stdout,
self.original_stderr: self.stderr,
sys.stdout: self.stdout,
sys.stderr: self.stderr,
}
restore_streams: dict[object, base.IO] = {}
if isinstance(self.stdout, WrappingIO):
restore_streams[self.stdout] = self.original_stdout
if isinstance(self.stderr, WrappingIO):
restore_streams[self.stderr] = self.original_stderr
seen: set[int] = set()
for logger_ in self._iter_loggers():
for handler in tuple(logger_.handlers):
if id(handler) in seen:
continue
seen.add(id(handler))
if not isinstance(handler, logging.StreamHandler):
continue
self._wrap_logging_handler(
handler,
wrapped_streams,
restore_streams,
)
def _wrap_logging_handler(
self,
handler: logging.StreamHandler[base.IO],
wrapped_streams: Mapping[types.Any, types.Any],
restore_streams: Mapping[types.Any, base.IO],
) -> None:
stream = handler.stream
replacement = wrapped_streams.get(stream)
if replacement is not None and replacement is not stream:
if self._set_handler_stream(handler, replacement):
self.logging_handlers.append((handler, stream))
elif (restore_stream := restore_streams.get(stream)) is not None:
self.logging_handlers.append((handler, restore_stream))
[docs]
def unwrap_logging(self) -> None:
if self.wrapped_logging > 1:
self.wrapped_logging -= 1
return
if not self.wrapped_logging:
return
while self.logging_handlers:
handler, stream = self.logging_handlers.pop()
self._set_handler_stream(handler, stream)
self.wrapped_logging = 0
def _set_handler_stream(
self,
handler: logging.StreamHandler[base.IO],
stream: types.Any,
) -> bool:
with contextlib.suppress(AttributeError, ValueError):
handler.setStream(stream)
return True
return False
def _iter_loggers(self) -> types.Iterator[logging.Logger]:
yield logging.getLogger()
for logger_ in tuple(logging.Logger.manager.loggerDict.values()):
if isinstance(logger_, logging.Logger):
yield logger_
[docs]
def unwrap_excepthook(self) -> None:
if self.wrapped_excepthook:
self.wrapped_excepthook -= 1
sys.excepthook = self.original_excepthook
[docs]
def wrap_excepthook(self) -> None:
if not self.wrapped_excepthook:
logger.debug('wrapping excepthook')
self.wrapped_excepthook += 1
sys.excepthook = self.excepthook
[docs]
def unwrap(self, stdout: bool = False, stderr: bool = False) -> None:
if stdout:
self.unwrap_stdout()
if stderr:
self.unwrap_stderr()
[docs]
def unwrap_stdout(self) -> None:
if self.wrapped_stdout > 1:
self.wrapped_stdout -= 1
else:
# Also reset our own reference so needs_clear() and
# update_capturing() don't act on a stale wrapper
self.stdout = sys.stdout = self.original_stdout
self.wrapped_stdout = 0
if not self.wrapped_stderr:
self.unwrap_excepthook()
[docs]
def unwrap_stderr(self) -> None:
if self.wrapped_stderr > 1:
self.wrapped_stderr -= 1
else:
# Also reset our own reference so needs_clear() and
# update_capturing() don't act on a stale wrapper
self.stderr = sys.stderr = self.original_stderr
self.wrapped_stderr = 0
if not self.wrapped_stdout:
self.unwrap_excepthook()
[docs]
def needs_clear(self) -> bool: # pragma: no cover
stdout_needs_clear = getattr(self.stdout, 'needs_clear', False)
stderr_needs_clear = getattr(self.stderr, 'needs_clear', False)
return stderr_needs_clear or stdout_needs_clear
[docs]
def flush(self) -> None:
if self.wrapped_stdout and isinstance(self.stdout, WrappingIO):
try:
self.stdout._flush()
except io.UnsupportedOperation:
self.wrapped_stdout = 0
logger.warning(
'Disabling stdout redirection, %r is not seekable',
sys.stdout,
)
if self.wrapped_stderr and isinstance(self.stderr, WrappingIO):
try:
self.stderr._flush()
except io.UnsupportedOperation: # pragma: no cover
self.wrapped_stderr = 0
logger.warning(
'Disabling stderr redirection, %r is not seekable',
sys.stderr,
)
[docs]
def excepthook(
self,
exc_type: type[BaseException],
exc_value: BaseException,
exc_traceback: TracebackType | None,
) -> None:
self.original_excepthook(exc_type, exc_value, exc_traceback)
self.flush()
[docs]
class AttributeDict(dict):
"""
A dict that can be accessed with .attribute.
>>> attrs = AttributeDict(spam=123)
# Reading
>>> attrs['spam']
123
>>> attrs.spam
123
# Read after update using attribute
>>> attrs.spam = 456
>>> attrs['spam']
456
>>> attrs.spam
456
# Read after update using dict access
>>> attrs['spam'] = 123
>>> attrs['spam']
123
>>> attrs.spam
123
# Read after update using dict access
>>> del attrs.spam
>>> attrs['spam']
Traceback (most recent call last):
...
KeyError: 'spam'
>>> attrs.spam
Traceback (most recent call last):
...
AttributeError: No such attribute: spam
>>> del attrs.spam
Traceback (most recent call last):
...
AttributeError: No such attribute: spam
"""
def __getattr__(self, name: str) -> typing.Any:
if name in self:
return self[name]
else:
raise AttributeError(f'No such attribute: {name}')
def __setattr__(self, name: str, value: typing.Any) -> None:
self[name] = value
def __delattr__(self, name: str) -> None:
if name in self:
del self[name]
else:
raise AttributeError(f'No such attribute: {name}')
streams = StreamWrapper()
atexit.register(streams.flush)