Implement __getitem__() protocol. Refs #35

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:

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':

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)
 

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):
     """

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"