From ceb4260484f4f30a0bb9f32d00c37f8363945e29 Mon Sep 17 00:00:00 2001 From: Timo Furrer Date: Wed, 4 Sep 2019 20:32:09 +0200 Subject: [PATCH 1/2] Implement __getitem__() protocol. Refs #35 --- README.md | 10 +++++ colorful/ansi.py | 3 ++ colorful/core.py | 76 +++++++++++++++++++++++++++++--- tests/test_core.py | 105 +++++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 187 insertions(+), 7 deletions(-) diff --git a/README.md b/README.md index 2b32293..5577029 100644 --- a/README.md +++ b/README.md @@ -362,6 +362,16 @@ print(colorful.red('red', nested=True) + ' default color') >>> assert len(s) == len(colorful.yellow(s)) ``` +#### Support the [`__getitem__()` protocol](https://docs.python.org/3/reference/datamodel.html#object.__getitem__) + +**colorful** tries to supports the `__getitem__()` protocol including [slices](https://docs.python.org/3/library/functions.html#slice) on the styled strings. + +However, there are some limitations in the current implementation: + +* Slices with negative steps are not supported +* All ANSI escape codes from the beginning of a string are included until the slice ends - even if they would cancel themselves out. +* The reset code (`\033[0m`) is always appended to slices and single index characters + ### Temporarily change colorful settings **colorful** provides a hand full of convenient context managers to change the colorful settings temporarily: diff --git a/colorful/ansi.py b/colorful/ansi.py index 47cc933..9331c32 100644 --- a/colorful/ansi.py +++ b/colorful/ansi.py @@ -47,6 +47,9 @@ #: Holds the base ANSI escape code ANSI_ESCAPE_CODE = '{csi}{{code}}m'.format(csi=CSI) +#: Holds the ANSI escape code to reset +ANSI_RESET_CODE = ANSI_ESCAPE_CODE.format(code=MODIFIERS["reset"][1]) + #: Holds the placeholder for the nest indicators NEST_PLACEHOLDER = ANSI_ESCAPE_CODE.format(code=26) diff --git a/colorful/core.py b/colorful/core.py index a2b053b..295ac29 100644 --- a/colorful/core.py +++ b/colorful/core.py @@ -10,17 +10,14 @@ :license: MIT, see LICENSE for more details. """ -from __future__ import print_function -from __future__ import unicode_literals +from __future__ import print_function, unicode_literals import os +import re import sys -from . import ansi -from . import colors -from . import styles -from . import terminal -from .utils import PY2, DEFAULT_ENCODING, UNICODE +from . import ansi, colors, styles, terminal +from .utils import DEFAULT_ENCODING, PY2, UNICODE #: Holds the name of the env variable which is # used as path to the default rgb.txt file @@ -313,6 +310,71 @@ def __getattr__(self, name): str_method = getattr(self.styled_string, name) return str_method + def __getitem__(self, item): + """Support indexing and slicing""" + if not isinstance(item, (int, slice)): + raise TypeError("ColorfulString indices must be integers") + + if isinstance(item, int): + start = item % len(self) + stop = start + 1 + step = 1 + else: + start, stop, step = item.indices(len(self)) + + if step < 0: + raise NotImplementedError("ColorfulString doesn't support negative slicing") + + sliced_orig_string = self.orig_string[item] + sliced_styled_string = "" + + #: Holds the regex pattern to match ANSI escape sequences which are not + # part of the slice + ansi_pattern = re.compile( + "^" + ansi.ANSI_ESCAPE_CODE.format(code=".*?").replace("[", r"\[") + ) + + #: Holds the index of the styled resp. the orig string while the slice + # is being consumed. + current_styled_string_idx = 0 + current_orig_string_idx = 0 + + #: Holds a counter to indicate how many steps of the ``step``-slice argument + # have to be made in order to consume the next char from the string. + step_counter = 0 + while current_styled_string_idx < len(self.styled_string): + ansi_match = ansi_pattern.search(self.styled_string[current_styled_string_idx:]) + if ansi_match: + # consume ANSI escape sequence from the string + sliced_styled_string += ansi_match.group(0) + + advance_style_idx = ansi_match.end(0) + advance_orig_idx = 0 + else: + if step_counter > 0: + # one-by-one consume the ``step``s. + step_counter -= 1 + else: + if start <= current_orig_string_idx < stop: + sliced_styled_string += self.orig_string[current_orig_string_idx] + step_counter = step - 1 + + advance_style_idx = 1 + advance_orig_idx = 1 + + # advance the counters of the styled and orig strings + # according to the consumed characters + current_styled_string_idx += advance_style_idx + current_orig_string_idx += advance_orig_idx + + if current_orig_string_idx >= stop: + # early exit if we discoved to be at the end of the slice + break + + # reset all colors and modifiers after the sliced string + sliced_styled_string += ansi.ANSI_RESET_CODE + return ColorfulString(sliced_orig_string, sliced_styled_string, self.colorful_ctx) + class Colorful(object): """ diff --git a/tests/test_core.py b/tests/test_core.py index 516b394..5fff0bb 100644 --- a/tests/test_core.py +++ b/tests/test_core.py @@ -872,3 +872,108 @@ def test_colorfulstyles_support_equals_protocol(style_a_name, style_b_name, expe # then assert actual_equal == expected_equal assert actual_hash_equal == expected_equal + + +def test_colorfulstring_only_support_int_and_slice_items(): + """Test that the Colorful __getitem__ protocol only supports ``int``s and ``slice``s""" + # given + colorful = core.Colorful(colormode=terminal.ANSI_8_COLORS) + s = colorful.red("Hello World") + + # then + with pytest.raises(TypeError, match="ColorfulString indices must be integers"): + # when + s["x"] + + +def test_colorfulstring_no_negative_slice_steps(): + """Test that the Colorful __getitem__ protocol doesn't support negative slicing steps""" + # given + colorful = core.Colorful(colormode=terminal.ANSI_8_COLORS) + s = colorful.red("Hello World") + + # then + with pytest.raises( + NotImplementedError, + match="ColorfulString doesn't support negative slicing" + ): + # when + s[0:1:-1] + + +def test_colorfulstring_get_char_at_positive_index(): + """Test getting single char from ColorfulString at a positive index""" + # given + colorful = core.Colorful(colormode=terminal.ANSI_8_COLORS) + s = colorful.red("Hello World") + + # when + sliced_s = s[4] + + # then + assert str(sliced_s) == "\033[31mo\033[0m" + + +def test_colorfulstring_get_char_at_negative_index(): + """Test getting single char from ColorfulString at a negative index""" + # given + colorful = core.Colorful(colormode=terminal.ANSI_8_COLORS) + s = colorful.red("Hello World") + + # when + sliced_s = s[-3] + + # then + assert str(sliced_s) == "\033[31mr\033[0m" + + +def test_colorfulstring_slice_string_with_single_color(): + """Test slicing a ColorfulString containing a single color""" + # given + colorful = core.Colorful(colormode=terminal.ANSI_8_COLORS) + s = colorful.red("Hello World") + + # when + sliced_s = s[4:7] + + # then + assert str(sliced_s) == "\033[31mo W\033[0m" + + +def test_colorfulstring_slice_with_step_in_string_with_single_color(): + """Test slicing a ColorfulString containing a single color""" + # given + colorful = core.Colorful(colormode=terminal.ANSI_8_COLORS) + s = colorful.red("Hello World") + + # when + sliced_s = s[2:9:3] + + # then + assert str(sliced_s) == "\033[31ml r\033[0m" + + +def test_colorfulstring_slice_string_with_two_colors(): + """Test slicing a ColorfulString consisting of two colors""" + # given + colorful = core.Colorful(colormode=terminal.ANSI_8_COLORS) + s = colorful.red("Hello") + " " + colorful.orange("World") + + # when + sliced_s = s[4:7] + + # then + assert str(sliced_s) == "\033[31mo\033[39m \033[33mW\033[0m" + + +def test_colorfulstring_slice_with_step_in_string_with_two_colors(): + """Test slicing a ColorfulString consisting of two colors""" + # given + colorful = core.Colorful(colormode=terminal.ANSI_8_COLORS) + s = colorful.red("Hello") + " " + colorful.orange("World") + + # when + sliced_s = s[2:9:3] + + # then + assert str(sliced_s) == "\033[31ml\033[39m \033[33mr\033[0m" From 326db2d71a4b94a08693bafd5d73f64441c8c67c Mon Sep 17 00:00:00 2001 From: Timo Furrer Date: Wed, 4 Sep 2019 20:33:21 +0200 Subject: [PATCH 2/2] alpha release: 0.6.0a1 --- colorful/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/colorful/__init__.py b/colorful/__init__.py index b45d377..ffdd129 100644 --- a/colorful/__init__.py +++ b/colorful/__init__.py @@ -21,7 +21,7 @@ from . import terminal #: Holds the current version -__version__ = '0.5.3' +__version__ = '0.6.0a1' # if we are on Windows we have to init colorama if platform.system() == 'Windows':