diff options
Diffstat (limited to 'tqdm/utils.py')
| -rw-r--r-- | tqdm/utils.py | 107 |
1 files changed, 76 insertions, 31 deletions
diff --git a/tqdm/utils.py b/tqdm/utils.py index 0632b8d..9883fda 100644 --- a/tqdm/utils.py +++ b/tqdm/utils.py @@ -4,31 +4,14 @@ General helpers required for `tqdm.std`. import os import re import sys -from functools import wraps +from functools import partial, partialmethod, wraps +from inspect import signature +# TODO consider using wcswidth third-party package for 0-width characters +from unicodedata import east_asian_width from warnings import warn from weakref import proxy -# py2/3 compat -try: - _range = xrange -except NameError: - _range = range - -try: - _unich = unichr -except NameError: - _unich = chr - -try: - _unicode = unicode -except NameError: - _unicode = str - -try: - _basestring = basestring -except NameError: - _basestring = str - +_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']) @@ -48,10 +31,78 @@ else: colorama.init() +def envwrap(prefix, types=None, is_method=False): + """ + Override parameter defaults via `os.environ[prefix + param_name]`. + Maps UPPER_CASE env vars map to lower_case param names. + camelCase isn't supported (because Windows ignores case). + + Precedence (highest first): + + - call (`foo(a=3)`) + - environ (`FOO_A=2`) + - signature (`def foo(a=1)`) + + Parameters + ---------- + prefix : str + Env var prefix, e.g. "FOO_" + types : dict, optional + Fallback mappings `{'param_name': type, ...}` if types cannot be + inferred from function signature. + Consider using `types=collections.defaultdict(lambda: ast.literal_eval)`. + is_method : bool, optional + Whether to use `functools.partialmethod`. If (default: False) use `functools.partial`. + + Examples + -------- + ``` + $ cat foo.py + from tqdm.utils import envwrap + @envwrap("FOO_") + def test(a=1, b=2, c=3): + print(f"received: a={a}, b={b}, c={c}") + + $ FOO_A=42 FOO_C=1337 python -c 'import foo; foo.test(c=99)' + received: a=42, b=2, c=99 + ``` + """ + if types is None: + types = {} + i = len(prefix) + env_overrides = {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 + # ignore unknown env vars + overrides = {k: v for k, v in env_overrides.items() if k in params} + # infer overrides' `type`s + for k in overrides: + param = params[k] + if param.annotation is not param.empty: # typehints + 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: # type of default value + overrides[k] = type(param.default)(overrides[k]) + else: + try: # `types` fallback + overrides[k] = types[k](overrides[k]) + except KeyError: # keep unconverted (`str`) + pass + return part(func, **overrides) + return wrap + + class FormatReplace(object): """ >>> a = FormatReplace('something') - >>> "{:5d}".format(a) + >>> f"{a:5d}" 'something' """ # NOQA: P102 def __init__(self, replace=''): @@ -320,14 +371,8 @@ def _term_move_up(): # pragma: no cover return '' if (os.name == 'nt') and (colorama is None) else '\x1b[A' -try: - # TODO consider using wcswidth third-party package for 0-width characters - from unicodedata import east_asian_width -except ImportError: - _text_width = len -else: - def _text_width(s): - return sum(2 if east_asian_width(ch) in 'FW' else 1 for ch in _unicode(s)) +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): |