From afe9ad895722cec2a7771b7b74aa0e4a03626bf1 Mon Sep 17 00:00:00 2001 From: Casper da Costa-Luis Date: Sun, 27 Sep 2020 19:26:13 +0100 Subject: [PATCH] support CLI colours - fixes #450 --- README.rst | 2 ++ tqdm/completion.sh | 4 ++-- tqdm/std.py | 56 +++++++++++++++++++++++++++++++++++++++------- tqdm/tqdm.1 | 9 +++++++- 4 files changed, 60 insertions(+), 11 deletions(-) diff --git a/README.rst b/README.rst index 51901f6bb..11624441f 100644 --- a/README.rst +++ b/README.rst @@ -460,6 +460,8 @@ Parameters The screen height. If specified, hides nested bars outside this bound. If unspecified, attempts to use environment height. The fallback is 20. +* colour : str or tuple, optional + Bar colour. May be RGB tuple: ``(0, 255, 0) == "GREEN"``. Extra CLI Options ~~~~~~~~~~~~~~~~~ diff --git a/tqdm/completion.sh b/tqdm/completion.sh index 42e84a8b5..788549cae 100755 --- a/tqdm/completion.sh +++ b/tqdm/completion.sh @@ -5,14 +5,14 @@ _tqdm(){ prv="${COMP_WORDS[COMP_CWORD - 1]}" case ${prv} in - --bar_format|--buf_size|--comppath|--delim|--desc|--initial|--lock_args|--manpath|--maxinterval|--mininterval|--miniters|--ncols|--nrows|--position|--postfix|--smoothing|--total|--unit|--unit_divisor) + --bar_format|--buf_size|--colour|--comppath|--delim|--desc|--initial|--lock_args|--manpath|--maxinterval|--mininterval|--miniters|--ncols|--nrows|--position|--postfix|--smoothing|--total|--unit|--unit_divisor) # await user input ;; "--log") COMPREPLY=($(compgen -W 'CRITICAL FATAL ERROR WARN WARNING INFO DEBUG NOTSET' -- ${cur})) ;; *) - COMPREPLY=($(compgen -W '--ascii --bar_format --buf_size --bytes --comppath --delim --desc --disable --dynamic_ncols --help --initial --leave --lock_args --log --manpath --maxinterval --mininterval --miniters --ncols --nrows --null --position --postfix --smoothing --tee --total --unit --unit_divisor --unit_scale --update --update_to --version --write_bytes -h -v' -- ${cur})) + COMPREPLY=($(compgen -W '--ascii --bar_format --buf_size --bytes --colour --comppath --delim --desc --disable --dynamic_ncols --help --initial --leave --lock_args --log --manpath --maxinterval --mininterval --miniters --ncols --nrows --null --position --postfix --smoothing --tee --total --unit --unit_divisor --unit_scale --update --update_to --version --write_bytes -h -v' -- ${cur})) ;; esac } diff --git a/tqdm/std.py b/tqdm/std.py index 7687c3703..b8a496a81 100644 --- a/tqdm/std.py +++ b/tqdm/std.py @@ -140,8 +140,13 @@ class Bar(object): ASCII = " 123456789#" UTF = u" " + u''.join(map(_unich, range(0x258F, 0x2587, -1))) BLANK = " " + COLOUR_RESET = '\x1b[0m' + COLOUR_RGB = '\x1b[38;2;%d;%d;%dm' + COLOURS = dict(BLACK='\x1b[30m', RED='\x1b[31m', GREEN='\x1b[32m', + YELLOW='\x1b[33m', BLUE='\x1b[34m', MAGENTA='\x1b[35m', + CYAN='\x1b[36m', WHITE='\x1b[37m') - def __init__(self, frac, default_len=10, charset=UTF): + def __init__(self, frac, default_len=10, charset=UTF, colour=None): if not (0 <= frac <= 1): warn("clamping frac to range [0, 1]", TqdmWarning, stacklevel=2) frac = max(0, min(1, frac)) @@ -149,6 +154,30 @@ def __init__(self, frac, default_len=10, charset=UTF): self.frac = frac self.default_len = default_len self.charset = charset + self.colour = colour + + @property + def colour(self): + return self._colour + + @colour.setter + def colour(self, value): + if not value: + self._colour = None + return + try: + if len(value) == 3: + self._colour = self.COLOUR_RGB % value + elif value.upper() in self.COLOURS: + self._colour = self.COLOURS[value.upper()] + else: + raise KeyError + except (KeyError, AttributeError): + warn(("Unknown colour (%s); valid choices:" + " [(0-255, 0-255, 0-255), %s]") % ( + value, ", ".join(self.COLOURS)), + TqdmWarning, stacklevel=2) + self._colour = None def __format__(self, format_spec): if format_spec: @@ -178,8 +207,10 @@ def __format__(self, format_spec): # whitespace padding if bar_length < N_BARS: - return bar + frac_bar + \ + bar = bar + frac_bar + \ charset[0] * (N_BARS - bar_length - 1) + if self.colour: + return self.colour + bar + self.COLOUR_RESET return bar @@ -309,7 +340,7 @@ def print_status(s): @staticmethod def format_meter(n, total, elapsed, ncols=None, prefix='', ascii=False, unit='it', unit_scale=False, rate=None, bar_format=None, - postfix=None, unit_divisor=1000, initial=0, + postfix=None, unit_divisor=1000, initial=0, colour=None, **extra_kwargs): """ Return a string-based progress bar given some parameters @@ -368,6 +399,8 @@ def format_meter(n, total, elapsed, ncols=None, prefix='', ascii=False, [default: 1000], ignored unless `unit_scale` is True. initial : int or float, optional The initial counter value [default: 0]. + colour : str or tuple, optional + Bar colour. May be RGB tuple: `(0, 255, 0) == "GREEN"`. Returns ------- @@ -442,6 +475,7 @@ def format_meter(n, total, elapsed, ncols=None, prefix='', ascii=False, rate_noinv_fmt=rate_noinv_fmt, rate_inv=inv_rate, rate_inv_fmt=rate_inv_fmt, postfix=postfix, unit_divisor=unit_divisor, + colour=colour, # plus more useful definitions remaining=remaining_str, remaining_s=remaining, l_bar=l_bar, r_bar=r_bar, @@ -483,7 +517,8 @@ def format_meter(n, total, elapsed, ncols=None, prefix='', ascii=False, frac, max(1, ncols - disp_len(nobar)) if ncols else 10, - charset=Bar.ASCII if ascii is True else ascii or Bar.UTF) + charset=Bar.ASCII if ascii is True else ascii or Bar.UTF, + colour=colour) if not _is_ascii(full_bar.charset) and _is_ascii(bar_format): bar_format = _unicode(bar_format) res = bar_format.format(bar=full_bar, **format_dict) @@ -501,7 +536,8 @@ def format_meter(n, total, elapsed, ncols=None, prefix='', ascii=False, 0, max(1, ncols - disp_len(nobar)) if ncols else 10, - charset=Bar.BLANK) + charset=Bar.BLANK, + colour=colour) res = bar_format.format(bar=full_bar, **format_dict) return disp_trim(res, ncols) if ncols else res else: @@ -802,7 +838,7 @@ def __init__(self, iterable=None, desc=None, total=None, leave=True, unit_scale=False, dynamic_ncols=False, smoothing=0.3, bar_format=None, initial=0, position=None, postfix=None, unit_divisor=1000, write_bytes=None, lock_args=None, - nrows=None, + nrows=None, colour=None, gui=False, **kwargs): """ Parameters @@ -908,6 +944,8 @@ def __init__(self, iterable=None, desc=None, total=None, leave=True, The screen height. If specified, hides nested bars outside this bound. If unspecified, attempts to use environment height. The fallback is 20. + colour : str or tuple, optional + Bar colour. May be RGB tuple: `(0, 255, 0) == "GREEN"`. gui : bool, optional WARNING: internal parameter - do not use. Use tqdm.gui.tqdm(...) instead. If set, will attempt to use @@ -1029,9 +1067,10 @@ def __init__(self, iterable=None, desc=None, total=None, leave=True, self.dynamic_ncols = dynamic_ncols self.smoothing = smoothing self.avg_time = None - self._time = time self.bar_format = bar_format self.postfix = None + self.colour = colour + self._time = time if postfix: try: self.set_postfix(refresh=False, **postfix) @@ -1453,7 +1492,8 @@ def format_dict(self): unit_scale=self.unit_scale, rate=1 / self.avg_time if self.avg_time else None, bar_format=self.bar_format, postfix=self.postfix, - unit_divisor=self.unit_divisor, initial=self.initial) + unit_divisor=self.unit_divisor, initial=self.initial, + colour=self.colour) def display(self, msg=None, pos=None): """ diff --git a/tqdm/tqdm.1 b/tqdm/tqdm.1 index e8dadade2..ab2e7c963 100644 --- a/tqdm/tqdm.1 +++ b/tqdm/tqdm.1 @@ -1,4 +1,4 @@ -.\" Automatically generated by Pandoc 1.19.2.1 +.\" Automatically generated by Pandoc 1.19.2 .\" .TH "TQDM" "1" "2015\-2020" "tqdm User Manuals" "" .hy @@ -227,6 +227,13 @@ The fallback is 20. .RS .RE .TP +.B \-\-colour=\f[I]colour\f[] +str or tuple, optional. +Bar colour. +May be RGB tuple: \f[C](0,\ 255,\ 0)\ ==\ "GREEN"\f[]. +.RS +.RE +.TP .B \-\-delim=\f[I]delim\f[] chr, optional. Delimiting character [default: \[aq]\\n\[aq]].