Complain if indices are out of bounds

MisanthropicBit · Feb 6, 2021 · 16330bb · 16330bb
1 parent 1bcf142
commit 16330bb
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 2 deletions.
diff --git a/src/colorise/__init__.py b/src/colorise/__init__.py
@@ -236,11 +236,14 @@ def highlight(
         return
 
     idx = 0
+    sorted_indices = sorted(indices)
+
+    if sorted_indices[0] < 0 or sorted_indices[-1] >= len(string):
+        raise IndexError('Index out of bounds')
 
     # Group consecutive indices, e.g. [0, 2, 3, 5, 6] -> [(0), (2, 3), (5, 6)]
-    # NOTE: The lambda syntax is necessary to support both Python 2 and 3
     groups = itertools.groupby(
-        enumerate(sorted(indices)),
+        enumerate(sorted_indices),
         lambda x: x[0] - x[1],
     )
 

diff --git a/tests/test_highlight.py b/tests/test_highlight.py
@@ -60,6 +60,18 @@ def test_invalid_highlight():
             colorise.highlight(text, indices, fg=color)
 
 
+def test_index_out_of_bounds():
+    text = 'This is a test'
+    index_tests = [
+        [0, 14, 3],
+        [0, -1, 3],
+    ]
+
+    for indices in index_tests:
+        with pytest.raises(IndexError):
+            colorise.highlight(text, indices, fg='red')
+
+
 @pytest.mark.skip_on_windows
 def test_highlight_named_output(test_stdout):
     test_stdout(