Refactor color_string()

scripts/wk/ui/ansi.py

@@ -3,7 +3,8 @@
 
 import itertools
 import logging
-import pathlib
+
+from typing import Iterable
 
 # STATIC VARIABLES
 LOG = logging.getLogger(__name__)
@@ -28,39 +29,36 @@ def clear_screen() -> None:
   print('\033c', end='', flush=True)
 
 
-def color_string(strings, colors, sep=' ') -> str:
+def color_string(
+    strings: Iterable[str] | str,
+    colors: Iterable[str | None] | str,
+    sep=' ',
+    ) -> str:
   """Build colored string using ANSI escapes, returns str."""
-  clear_code = COLORS['CLEAR']
+  data = {'strings': strings, 'colors': colors}
   msg = []
 
-  # Convert to tuples if necessary
+  # Convert input to tuples of strings
-  if isinstance(strings, (str, pathlib.Path)):
+  for k, v in data.items():
-    strings = (strings,)
+    if isinstance(v, str):
-  if isinstance(colors, (str, pathlib.Path)):
+      # Avoid splitting string into a list of characters
-    colors = (colors,)
+      data[k] = (v,)
-
+    try:
-  # Convert to strings if necessary
+      iter(v)
-  try:
+    except TypeError:
-    iter(strings)
+      # Assuming single element passed, convert to string
-  except TypeError:
+      data[k] = (str(v),)
-    # Assuming single element passed, convert to string
-    strings = (str(strings),)
-  try:
-    iter(colors)
-  except TypeError:
-    # Assuming single element passed, convert to string
-    colors = (str(colors),)
 
   # Build new string with color escapes added
-  for string, color in itertools.zip_longest(strings, colors):
+  for string, color in itertools.zip_longest(data['strings'], data['colors']):
-    color_code = COLORS.get(color, clear_code)
+    color_code = COLORS.get(str(color), COLORS['CLEAR'])
-    msg.append(f'{color_code}{string}{clear_code}')
+    msg.append(f'{color_code}{string}{COLORS["CLEAR"]}')
 
   # Done
   return sep.join(msg)
 
 
-def strip_colors(string) -> str:
+def strip_colors(string: str) -> str:
   """Strip known ANSI color escapes from string, returns str."""
   LOG.debug('string: %s', string)
   for color in COLORS.values():