codekingpro/portable-devtools

1"""2General helpers required for `tqdm.std`.3"""4import os5import re6import sys7from functools import partial, partialmethod, wraps8from inspect import signature9# TODO consider using wcswidth third-party package for 0-width characters10from unicodedata import east_asian_width11from warnings import warn12from weakref import proxy13 14_range, _unich, _unicode, _basestring = range, chr, str, str15CUR_OS = sys.platform16IS_WIN = any(CUR_OS.startswith(i) for i in ['win32', 'cygwin'])17IS_NIX = any(CUR_OS.startswith(i) for i in ['aix', 'linux', 'darwin', 'freebsd'])18RE_ANSI = re.compile(r"\x1b\[[;\d]*[A-Za-z]")19 20try:21    if IS_WIN:22        import colorama23    else:24        raise ImportError25except ImportError:26    colorama = None27else:28    try:29        colorama.init(strip=False)30    except TypeError:31        colorama.init()32 33 34def envwrap(prefix, types=None, is_method=False):35    """36    Override parameter defaults via `os.environ[prefix + param_name]`.37    Maps UPPER_CASE env vars map to lower_case param names.38    camelCase isn't supported (because Windows ignores case).39 40    Precedence (highest first):41 42    - call (`foo(a=3)`)43    - environ (`FOO_A=2`)44    - signature (`def foo(a=1)`)45 46    Parameters47    ----------48    prefix  : str49        Env var prefix, e.g. "FOO_"50    types  : dict, optional51        Fallback mappings `{'param_name': type, ...}` if types cannot be52        inferred from function signature.53        Consider using `types=collections.defaultdict(lambda: ast.literal_eval)`.54    is_method  : bool, optional55        Whether to use `functools.partialmethod`. If (default: False) use `functools.partial`.56 57    Examples58    --------59    ```60    $ cat foo.py61    from tqdm.utils import envwrap62    @envwrap("FOO_")63    def test(a=1, b=2, c=3):64        print(f"received: a={a}, b={b}, c={c}")65 66    $ FOO_A=42 FOO_C=1337 python -c 'import foo; foo.test(c=99)'67    received: a=42, b=2, c=9968    ```69    """70    if types is None:71        types = {}72    i = len(prefix)73    env_overrides = {k[i:].lower(): v for k, v in os.environ.items() if k.startswith(prefix)}74    part = partialmethod if is_method else partial75 76    def wrap(func):77        params = signature(func).parameters78        # ignore unknown env vars79        overrides = {k: v for k, v in env_overrides.items() if k in params}80        # infer overrides' `type`s81        for k in overrides:82            param = params[k]83            if param.annotation is not param.empty:  # typehints84                for typ in getattr(param.annotation, '__args__', (param.annotation,)):85                    try:86                        overrides[k] = typ(overrides[k])87                    except Exception:88                        pass89                    else:90                        break91            elif param.default is not None:  # type of default value92                overrides[k] = type(param.default)(overrides[k])93            else:94                try:  # `types` fallback95                    overrides[k] = types[k](overrides[k])96                except KeyError:  # keep unconverted (`str`)97                    pass98        return part(func, **overrides)99    return wrap100 101 102class FormatReplace:103    """104    >>> a = FormatReplace('something')105    >>> f"{a:5d}"106    'something'107    """  # NOQA: P102108    def __init__(self, replace=''):109        self.replace = replace110        self.format_called = 0111 112    def __format__(self, _):113        self.format_called += 1114        return self.replace115 116 117class Comparable:118    """Assumes child has self._comparable attr/@property"""119    def __lt__(self, other):120        return self._comparable < other._comparable121 122    def __le__(self, other):123        return (self < other) or (self == other)124 125    def __eq__(self, other):126        return self._comparable == other._comparable127 128    def __ne__(self, other):129        return not self == other130 131    def __gt__(self, other):132        return not self <= other133 134    def __ge__(self, other):135        return not self < other136 137 138class ObjectWrapper:139    def __getattr__(self, name):140        return getattr(self._wrapped, name)141 142    def __setattr__(self, name, value):143        return setattr(self._wrapped, name, value)144 145    def wrapper_getattr(self, name):146        """Actual `self.getattr` rather than self._wrapped.getattr"""147        try:148            return object.__getattr__(self, name)149        except AttributeError:  # py2150            return getattr(self, name)151 152    def wrapper_setattr(self, name, value):153        """Actual `self.setattr` rather than self._wrapped.setattr"""154        return object.__setattr__(self, name, value)155 156    def __init__(self, wrapped):157        """158        Thin wrapper around a given object159        """160        self.wrapper_setattr('_wrapped', wrapped)161 162 163class SimpleTextIOWrapper(ObjectWrapper):164    """165    Change only `.write()` of the wrapped object by encoding the passed166    value and passing the result to the wrapped object's `.write()` method.167    """168    # pylint: disable=too-few-public-methods169    def __init__(self, wrapped, encoding):170        super().__init__(wrapped)171        self.wrapper_setattr('encoding', encoding)172 173    def write(self, s):174        """175        Encode `s` and pass to the wrapped object's `.write()` method.176        """177        return self._wrapped.write(s.encode(self.wrapper_getattr('encoding')))178 179    def __eq__(self, other):180        return self._wrapped == getattr(other, '_wrapped', other)181 182 183class DisableOnWriteError(ObjectWrapper):184    """185    Disable the given `tqdm_instance` upon `write()` or `flush()` errors.186    """187    @staticmethod188    def disable_on_exception(tqdm_instance, func):189        """190        Quietly set `tqdm_instance.miniters=inf` if `func` raises `errno=5`.191        """192        tqdm_instance = proxy(tqdm_instance)193 194        def inner(*args, **kwargs):195            try:196                return func(*args, **kwargs)197            except OSError as e:198                if e.errno != 5:199                    raise200                try:201                    tqdm_instance.miniters = float('inf')202                except ReferenceError:203                    pass204            except ValueError as e:205                if 'closed' not in str(e):206                    raise207                try:208                    tqdm_instance.miniters = float('inf')209                except ReferenceError:210                    pass211        return inner212 213    def __init__(self, wrapped, tqdm_instance):  # noqa: B042214        super().__init__(wrapped)215        if hasattr(wrapped, 'write'):216            self.wrapper_setattr(217                'write', self.disable_on_exception(tqdm_instance, wrapped.write))218        if hasattr(wrapped, 'flush'):219            self.wrapper_setattr(220                'flush', self.disable_on_exception(tqdm_instance, wrapped.flush))221 222    def __eq__(self, other):223        return self._wrapped == getattr(other, '_wrapped', other)224 225 226class CallbackIOWrapper(ObjectWrapper):227    def __init__(self, callback, stream, method="read"):228        """229        Wrap a given `file`-like object's `read()` or `write()` to report230        lengths to the given `callback`231        """232        super().__init__(stream)233        func = getattr(stream, method)234        if method == "write":235            @wraps(func)236            def write(data, *args, **kwargs):237                res = func(data, *args, **kwargs)238                callback(len(data))239                return res240            self.wrapper_setattr('write', write)241        elif method == "read":242            @wraps(func)243            def read(*args, **kwargs):244                data = func(*args, **kwargs)245                callback(len(data))246                return data247            self.wrapper_setattr('read', read)248        else:249            raise KeyError("Can only wrap read/write methods")250 251 252def _is_utf(encoding):253    try:254        '\u2588\u2589'.encode(encoding)255    except UnicodeEncodeError:256        return False257    except Exception:258        try:259            return encoding.lower().startswith('utf-') or ('U8' == encoding)260        except Exception:261            return False262    else:263        return True264 265 266def _supports_unicode(fp):267    try:268        return _is_utf(fp.encoding)269    except AttributeError:270        return False271 272 273def _is_ascii(s):274    if isinstance(s, str):275        for c in s:276            if ord(c) > 255:277                return False278        return True279    return _supports_unicode(s)280 281 282def _screen_shape_wrapper():  # pragma: no cover283    """284    Return a function which returns console dimensions (width, height).285    Supported: linux, osx, windows, cygwin.286    """287    _screen_shape = None288    if IS_WIN:289        _screen_shape = _screen_shape_windows290        if _screen_shape is None:291            _screen_shape = _screen_shape_tput292    if IS_NIX:293        _screen_shape = _screen_shape_linux294    return _screen_shape295 296 297def _screen_shape_windows(fp):  # pragma: no cover298    try:299        import struct300        from ctypes import create_string_buffer, windll301        from sys import stdin, stdout302 303        io_handle = -12  # assume stderr304        if fp == stdin:305            io_handle = -10306        elif fp == stdout:307            io_handle = -11308 309        h = windll.kernel32.GetStdHandle(io_handle)310        csbi = create_string_buffer(22)311        res = windll.kernel32.GetConsoleScreenBufferInfo(h, csbi)312        if res:313            (_bufx, _bufy, _curx, _cury, _wattr, left, top, right, bottom,314             _maxx, _maxy) = struct.unpack("hhhhHhhhhhh", csbi.raw)315            return right - left, bottom - top  # +1316    except Exception:  # nosec317        pass318    return None, None319 320 321def _screen_shape_tput(*_):  # pragma: no cover322    """cygwin xterm (windows)"""323    try:324        import shlex325        from subprocess import check_call  # nosec326        return [int(check_call(shlex.split('tput ' + i))) - 1327                for i in ('cols', 'lines')]328    except Exception:  # nosec329        pass330    return None, None331 332 333def _screen_shape_linux(fp):  # pragma: no cover334 335    try:336        from array import array337        from fcntl import ioctl338        from termios import TIOCGWINSZ339    except ImportError:340        return None, None341    else:342        try:343            rows, cols = array('h', ioctl(fp, TIOCGWINSZ, '\0' * 8))[:2]344            return cols, rows345        except Exception:346            try:347                return [int(os.environ[i]) - 1 for i in ("COLUMNS", "LINES")]348            except (KeyError, ValueError):349                return None, None350 351 352def _environ_cols_wrapper():  # pragma: no cover353    """354    Return a function which returns console width.355    Supported: linux, osx, windows, cygwin.356    """357    warn("Use `_screen_shape_wrapper()(file)[0]` instead of"358         " `_environ_cols_wrapper()(file)`", DeprecationWarning, stacklevel=2)359    shape = _screen_shape_wrapper()360    if not shape:361        return None362 363    @wraps(shape)364    def inner(fp):365        return shape(fp)[0]366 367    return inner368 369 370def _term_move_up():  # pragma: no cover371    return '' if (os.name == 'nt') and (colorama is None) else '\x1b[A'372 373 374def _text_width(s):375    return sum(2 if east_asian_width(ch) in 'FW' else 1 for ch in str(s))376 377 378def disp_len(data):379    """380    Returns the real on-screen length of a string which may contain381    ANSI control codes and wide chars.382    """383    return _text_width(RE_ANSI.sub('', data))384 385 386def disp_trim(data, length):387    """388    Trim a string which may contain ANSI control characters.389    """390    if len(data) == disp_len(data):391        return data[:length]392 393    ansi_present = bool(RE_ANSI.search(data))394    while disp_len(data) > length:  # carefully delete one char at a time395        data = data[:-1]396    if ansi_present and bool(RE_ANSI.search(data)):397        # assume ANSI reset is required398        return data if data.endswith("\033[0m") else data + "\033[0m"399    return data400