Made len_color() safer and added testing

wolph · Jul 13, 2019 · 34aff68 · 34aff68
1 parent 4b60e82
commit 34aff68
Showing 1 changed file with 19 additions and 3 deletions.
diff --git a/progressbar/utils.py b/progressbar/utils.py
@@ -19,9 +19,25 @@
 
 
 def len_color(value):
-    '''Return the length of `value` without ANSI escape codes'''
-    if isinstance(value, str):
-        value = re.sub(u'\u001b\\[.*?[@-~]', '', value)
+    '''
+    Return the length of `value` without ANSI escape codes
+
+    >>> len_color(u'\u001b[1234]abc')
+    3
+    >>> len_color(b'\u001b[1234]abc')
+    3
+    >>> len_color('\u001b[1234]abc')
+    3
+    '''
+    pattern = u'\u001b\\[.*?[@-~]'
+    if isinstance(value, bytes):
+        pattern = pattern.encode()
+        replace = b''
+        assert isinstance(pattern, bytes)
+    else:
+        replace = ''
+
+    value = re.sub(pattern, replace, value)
     return len(value)