| """ |
| General helpers required for `tqdm.std`. |
| """ |
| import os |
| import re |
| import sys |
| from functools import partial, partialmethod, wraps |
| from inspect import signature |
| |
| from unicodedata import east_asian_width |
| from warnings import warn |
| from weakref import proxy |
|
|
| _range, _unich, _unicode, _basestring = range, chr, str, str |
| CUR_OS = sys.platform |
| IS_WIN = any(CUR_OS.startswith(i) for i in ['win32', 'cygwin']) |
| IS_NIX = any(CUR_OS.startswith(i) for i in ['aix', 'linux', 'darwin', 'freebsd']) |
| RE_ANSI = re.compile(r"\x1b\[[;\d]*[A-Za-z]") |
|
|
| try: |
| if IS_WIN: |
| import colorama |
| else: |
| raise ImportError |
| except ImportError: |
| colorama = None |
| else: |
| try: |
| colorama.init(strip=False) |
| except TypeError: |
| colorama.init() |
|
|
|
|
| def envwrap(name, app="", types=None, is_method=False): |
| """ |
| Basic (env-only) version of [envwrap](https://github.com/tqdm/envwrap). |
| Install `envwrap` for config file support. |
| """ |
| if types is None: |
| types = {} |
| if name[-1] == "_": |
| name = name[:-1] |
| warn("Trailing underscore in `name` is automatic", DeprecationWarning, stacklevel=2) |
| prefixes = (name, f"{name}_{app}") if app else (name,) |
| env_overrides = {} |
| for prefix in prefixes: |
| prefix = prefix.upper() + "_" |
| i = len(prefix) |
| env_overrides.update( |
| (k[i:].lower(), v) for k, v in os.environ.items() if k.startswith(prefix)) |
| part = partialmethod if is_method else partial |
|
|
| def wrap(func): |
| params = signature(func).parameters |
| |
| overrides = {k: v for k, v in env_overrides.items() if k in params} |
| |
| for k in overrides: |
| param = params[k] |
| if param.annotation is not param.empty: |
| for typ in getattr(param.annotation, '__args__', (param.annotation,)): |
| try: |
| overrides[k] = typ(overrides[k]) |
| except Exception: |
| pass |
| else: |
| break |
| elif param.default is not None: |
| overrides[k] = type(param.default)(overrides[k]) |
| else: |
| try: |
| overrides[k] = types[k](overrides[k]) |
| except KeyError: |
| pass |
| return part(func, **overrides) |
| return wrap |
|
|
|
|
| try: |
| from envwrap import envwrap |
| except ModuleNotFoundError: |
| pass |
|
|
|
|
| class FormatReplace: |
| """ |
| >>> a = FormatReplace('something') |
| >>> f"{a:5d}" |
| 'something' |
| """ |
| def __init__(self, replace=''): |
| self.replace = replace |
| self.format_called = 0 |
|
|
| def __format__(self, _): |
| self.format_called += 1 |
| return self.replace |
|
|
|
|
| class Comparable: |
| """Assumes child has self._comparable attr/@property""" |
| def __lt__(self, other): |
| return self._comparable < other._comparable |
|
|
| def __le__(self, other): |
| return (self < other) or (self == other) |
|
|
| def __eq__(self, other): |
| return self._comparable == other._comparable |
|
|
| def __ne__(self, other): |
| return not self == other |
|
|
| def __gt__(self, other): |
| return not self <= other |
|
|
| def __ge__(self, other): |
| return not self < other |
|
|
|
|
| class ObjectWrapper: |
| def __getattr__(self, name): |
| return getattr(self._wrapped, name) |
|
|
| def __setattr__(self, name, value): |
| return setattr(self._wrapped, name, value) |
|
|
| def wrapper_getattr(self, name): |
| """Actual `self.getattr` rather than self._wrapped.getattr""" |
| try: |
| return object.__getattr__(self, name) |
| except AttributeError: |
| return getattr(self, name) |
|
|
| def wrapper_setattr(self, name, value): |
| """Actual `self.setattr` rather than self._wrapped.setattr""" |
| return object.__setattr__(self, name, value) |
|
|
| def __init__(self, wrapped): |
| """ |
| Thin wrapper around a given object |
| """ |
| self.wrapper_setattr('_wrapped', wrapped) |
|
|
|
|
| class SimpleTextIOWrapper(ObjectWrapper): |
| """ |
| Change only `.write()` of the wrapped object by encoding the passed |
| value and passing the result to the wrapped object's `.write()` method. |
| """ |
| |
| def __init__(self, wrapped, encoding): |
| super().__init__(wrapped) |
| self.wrapper_setattr('encoding', encoding) |
|
|
| def write(self, s): |
| """ |
| Encode `s` and pass to the wrapped object's `.write()` method. |
| """ |
| return self._wrapped.write(s.encode(self.wrapper_getattr('encoding'))) |
|
|
| def __eq__(self, other): |
| return self._wrapped == getattr(other, '_wrapped', other) |
|
|
|
|
| class DisableOnWriteError(ObjectWrapper): |
| """ |
| Disable the given `tqdm_instance` upon `write()` or `flush()` errors. |
| """ |
| @staticmethod |
| def disable_on_exception(tqdm_instance, func): |
| """ |
| Quietly set `tqdm_instance.miniters=inf` if `func` raises `errno=5`. |
| """ |
| tqdm_instance = proxy(tqdm_instance) |
|
|
| def inner(*args, **kwargs): |
| try: |
| return func(*args, **kwargs) |
| except OSError as e: |
| if e.errno != 5: |
| raise |
| try: |
| tqdm_instance.miniters = float('inf') |
| except ReferenceError: |
| pass |
| except ValueError as e: |
| if 'closed' not in str(e): |
| raise |
| try: |
| tqdm_instance.miniters = float('inf') |
| except ReferenceError: |
| pass |
| return inner |
|
|
| def __init__(self, wrapped, tqdm_instance): |
| super().__init__(wrapped) |
| if hasattr(wrapped, 'write'): |
| self.wrapper_setattr( |
| 'write', self.disable_on_exception(tqdm_instance, wrapped.write)) |
| if hasattr(wrapped, 'flush'): |
| self.wrapper_setattr( |
| 'flush', self.disable_on_exception(tqdm_instance, wrapped.flush)) |
|
|
| def __eq__(self, other): |
| return self._wrapped == getattr(other, '_wrapped', other) |
|
|
|
|
| class CallbackIOWrapper(ObjectWrapper): |
| def __init__(self, callback, stream, method="read"): |
| """ |
| Wrap a given `file`-like object's `read()` or `write()` to report |
| lengths to the given `callback` |
| """ |
| super().__init__(stream) |
| func = getattr(stream, method) |
| if method == "write": |
| @wraps(func) |
| def write(data, *args, **kwargs): |
| res = func(data, *args, **kwargs) |
| callback(len(data)) |
| return res |
| self.wrapper_setattr('write', write) |
| elif method == "read": |
| @wraps(func) |
| def read(*args, **kwargs): |
| data = func(*args, **kwargs) |
| callback(len(data)) |
| return data |
| self.wrapper_setattr('read', read) |
| else: |
| raise KeyError("Can only wrap read/write methods") |
|
|
|
|
| def _is_utf(encoding): |
| try: |
| '\u2588\u2589'.encode(encoding) |
| except UnicodeEncodeError: |
| return False |
| except Exception: |
| try: |
| return encoding.lower().startswith('utf-') or ('U8' == encoding) |
| except Exception: |
| return False |
| else: |
| return True |
|
|
|
|
| def _supports_unicode(fp): |
| try: |
| return _is_utf(fp.encoding) |
| except AttributeError: |
| return False |
|
|
|
|
| def _is_ascii(s): |
| if isinstance(s, str): |
| for c in s: |
| if ord(c) > 255: |
| return False |
| return True |
| return _supports_unicode(s) |
|
|
|
|
| def _screen_shape_wrapper(): |
| """ |
| Return a function which returns console dimensions (width, height). |
| Supported: linux, osx, windows, cygwin. |
| """ |
| def inner(fp): |
| try: |
| from os import get_terminal_size |
| cols, lines = get_terminal_size(getattr(fp, 'fileno', lambda: None)()) |
| return cols - 1, lines - 1 |
| except Exception: |
| return None, None |
|
|
| return inner |
|
|
|
|
| def _environ_cols_wrapper(): |
| """ |
| Return a function which returns console width. |
| Supported: linux, osx, windows, cygwin. |
| """ |
| warn("Use `_screen_shape_wrapper()(file)[0]` instead of" |
| " `_environ_cols_wrapper()(file)`", DeprecationWarning, stacklevel=2) |
| shape = _screen_shape_wrapper() |
| if not shape: |
| return None |
|
|
| @wraps(shape) |
| def inner(fp): |
| return shape(fp)[0] |
|
|
| return inner |
|
|
|
|
| def _term_move_up(): |
| return '' if (os.name == 'nt') and (colorama is None) else '\x1b[A' |
|
|
|
|
| def _text_width(s): |
| return sum(2 if east_asian_width(ch) in 'FW' else 1 for ch in str(s)) |
|
|
|
|
| def disp_len(data): |
| """ |
| Returns the real on-screen length of a string which may contain |
| ANSI control codes and wide chars. |
| """ |
| return _text_width(RE_ANSI.sub('', data)) |
|
|
|
|
| def disp_trim(data, length): |
| """ |
| Trim a string which may contain ANSI control characters. |
| """ |
| if len(data) == disp_len(data): |
| return data[:length] |
|
|
| ansi_present = bool(RE_ANSI.search(data)) |
| while disp_len(data) > length: |
| data = data[:-1] |
| if ansi_present and bool(RE_ANSI.search(data)): |
| |
| return data if data.endswith("\033[0m") else data + "\033[0m" |
| return data |
|
|