added color support for rotating markers

wolph · wolph · commit 9616cee18bfe · 2020-02-25T16:48:15.000+01:00
diff --git a/progressbar/widgets.py b/progressbar/widgets.py
@@ -9,6 +9,7 @@
 import sys
 import pprint
 import datetime
+import functools
 
 from python_utils import converters
 
@@ -32,7 +33,49 @@ def render_input(progress, data, width):
         return input_
 
 
-def create_marker(marker):
+def create_wrapper(wrapper):
+    '''Convert a wrapper tuple or format string to a format string
+
+    >>> create_wrapper('')
+
+    >>> create_wrapper('a{}b')
+    'a{}b'
+
+    >>> create_wrapper(('a', 'b'))
+    'a{}b'
+    '''
+    if isinstance(wrapper, tuple) and len(wrapper) == 2:
+        a, b = wrapper
+        wrapper = (a or '') + '{}' + (b or '')
+    elif not wrapper:
+        return
+
+    if isinstance(wrapper, six.string_types):
+        assert '{}' in wrapper, 'Expected string with {} for formatting'
+    else:
+        raise RuntimeError('Pass either a begin/end string as a tuple or a'
+                           ' template string with {}')
+
+    return wrapper
+
+
+def wrapper(function, wrapper):
+    '''Wrap the output of a function in a template string or a tuple with
+    begin/end strings
+
+    '''
+    wrapper = create_wrapper(wrapper)
+    if not wrapper:
+        return function
+
+    @functools.wraps(function)
+    def wrap(*args, **kwargs):
+        return wrapper.format(function(*args, **kwargs))
+
+    return wrap
+
+
+def create_marker(marker, wrap=None):
     def _marker(progress, data, width):
         if progress.max_value is not base.UnknownLength \
                 and progress.max_value > 0:
@@ -45,9 +88,9 @@ def _marker(progress, data, width):
         marker = converters.to_unicode(marker)
         assert utils.len_color(marker) == 1, \
             'Markers are required to be 1 char'
-        return _marker
+        return wrapper(_marker, wrap)
     else:
-        return marker
+        return wrapper(marker, wrap)
 
 
 class FormatWidgetMixin(object):
@@ -525,10 +568,13 @@ class AnimatedMarker(TimeSensitiveWidgetBase):
     it were rotating.
     '''
 
-    def __init__(self, markers='|/-\\', default=None, fill='', **kwargs):
+    def __init__(self, markers='|/-\\', default=None, fill='',
+                 marker_wrap=None, fill_wrap=None, **kwargs):
         self.markers = markers
+        self.marker_wrap = create_wrapper(marker_wrap)
         self.default = default or markers[0]
-        self.fill = create_marker(fill) if fill else None
+        self.fill_wrap = create_wrapper(fill_wrap)
+        self.fill = create_marker(fill, self.fill_wrap) if fill else None
         WidgetBase.__init__(self, **kwargs)
 
     def __call__(self, progress, data, width=None):
@@ -538,14 +584,17 @@ def __call__(self, progress, data, width=None):
         if progress.end_time:
             return self.default
 
+        marker = self.markers[data['updates'] % len(self.markers)]
+        if self.marker_wrap:
+            marker = self.marker_wrap.format(marker)
+
         if self.fill:
             # Cut the last character so we can replace it with our marker
-            fill = self.fill(progress, data, width)[:-1]
+            fill = self.fill(progress, data, width - progress.custom_len(
+                marker))
         else:
             fill = ''
-
-        marker = self.markers[data['updates'] % len(self.markers)]
-
+        
         # Python 3 returns an int when indexing bytes
         if isinstance(marker, int):  # pragma: no cover
             marker = bytes(marker)
@@ -642,7 +691,7 @@ class Bar(AutoWidthWidgetBase):
     '''A progress bar which stretches to fill the line.'''
 
     def __init__(self, marker='#', left='|', right='|', fill=' ',
-                 fill_left=True, **kwargs):
+                 fill_left=True, marker_wrap=None, **kwargs):
         '''Creates a customizable progress bar.
 
         The callable takes the same parameters as the `__call__` method
@@ -654,7 +703,7 @@ def __init__(self, marker='#', left='|', right='|', fill=' ',
         fill_left - whether to fill from the left or the right
         '''
 
-        self.marker = create_marker(marker)
+        self.marker = create_marker(marker, marker_wrap)
         self.left = string_or_lambda(left)
         self.right = string_or_lambda(right)
         self.fill = string_or_lambda(fill)
@@ -670,6 +719,10 @@ def __call__(self, progress, data, width):
         width -= progress.custom_len(left) + progress.custom_len(right)
         marker = converters.to_unicode(self.marker(progress, data, width))
         fill = converters.to_unicode(self.fill(progress, data, width))
+
+        # Make sure we ignore invisible characters when filling
+        width += len(marker) - progress.custom_len(marker)
+
         if self.fill_left:
             marker = marker.ljust(width, fill)
         else:
@@ -796,7 +849,7 @@ def __call__(self, progress, data, width):
             width_accumulated = 0
             for marker, value in zip(self.markers, values):
                 marker = converters.to_unicode(marker(progress, data, width))
-                assert utils.len_color(marker) == 1
+                assert progress.custom_len(marker) == 1
 
                 values_accumulated += value
                 item_width = int(values_accumulated / values_sum * width)
@@ -805,7 +858,7 @@ def __call__(self, progress, data, width):
                 middle += item_width * marker
         else:
             fill = converters.to_unicode(self.fill(progress, data, width))
-            assert utils.len_color(fill) == 1
+            assert progress.custom_len(fill) == 1
             middle = fill * width
 
         return left + middle + right
diff --git a/tests/test_widgets.py b/tests/test_widgets.py
@@ -6,6 +6,14 @@
 max_values = [None, 10, progressbar.UnknownLength]
 
 
+def test_create_wrapper():
+    with pytest.raises(AssertionError):
+        progressbar.widgets.create_wrapper('ab')
+
+    with pytest.raises(RuntimeError):
+        progressbar.widgets.create_wrapper(123)
+
+
 def test_widgets_small_values():
     widgets = [
         'Test: ',
