1 files changed, 49 insertions, 0 deletions
diff --git a/colorclass/search.py b/colorclass/search.py
new file mode 100644
index 0000000..555402d
--- /dev/null
+++ b/colorclass/search.py
@@ -0,0 +1,49 @@
+"""Determine color of characters that may or may not be adjacent to ANSI escape sequences."""
+
+from colorclass.parse import RE_SPLIT
+
+
+def build_color_index(ansi_string):
+    """Build an index between visible characters and a string with invisible color codes.
+
+    :param str ansi_string: String with color codes (ANSI escape sequences).
+
+    :return: Position of visible characters in color string (indexes match non-color string).
+    :rtype: tuple
+    """
+    mapping = list()
+    color_offset = 0
+    for item in (i for i in RE_SPLIT.split(ansi_string) if i):
+        if RE_SPLIT.match(item):
+            color_offset += len(item)
+        else:
+            for _ in range(len(item)):
+                mapping.append(color_offset)
+                color_offset += 1
+    return tuple(mapping)
+
+
+def find_char_color(ansi_string, pos):
+    """Determine what color a character is in the string.
+
+    :param str ansi_string: String with color codes (ANSI escape sequences).
+    :param int pos: Position of the character in the ansi_string.
+
+    :return: Character along with all surrounding color codes.
+    :rtype: str
+    """
+    result = list()
+    position = 0  # Set to None when character is found.
+    for item in (i for i in RE_SPLIT.split(ansi_string) if i):
+        if RE_SPLIT.match(item):
+            result.append(item)
+            if position is not None:
+                position += len(item)
+        elif position is not None:
+            for char in item:
+                if position == pos:
+                    result.append(char)
+                    position = None
+                    break
+                position += 1
+    return ''.join(result)