Improve coloring

coconut/terminal.py

@@ -183,6 +183,16 @@ def logging(self):
                 sys.stdout = old_stdout
 
 
+def should_use_color(file=None):
+    """Determine if colors should be used for the given file object."""
+    use_color = get_bool_env_var(use_color_env_var, default=None)
+    if use_color is not None:
+        return use_color
+    if get_bool_env_var("CLICOLOR_FORCE") or get_bool_env_var("FORCE_COLOR"):
+        return True
+    return file is not None and not isatty(file)
+
+
 # -----------------------------------------------------------------------------------------------------------------------
 # LOGGER:
 # -----------------------------------------------------------------------------------------------------------------------
@@ -210,11 +220,7 @@ def __init__(self, other=None):
     @classmethod
     def enable_colors(cls, file=None):
         """Attempt to enable CLI colors."""
-        use_color = get_bool_env_var(use_color_env_var, default=None)
-        if (
-            use_color is False
-            or use_color is None and file is not None and not isatty(file)
-        ):
+        if not should_use_color(file):
             return False
         if not cls.colors_enabled:
             # necessary to resolve https://bugs.python.org/issue40134

coconut/util.py

@@ -8,7 +8,7 @@
 """
 Author: Evan Hubinger
 License: Apache 2.0
-Description: Installer for the Coconut Jupyter kernel.
+Description: Base Coconut utilities.
 """
 
 # -----------------------------------------------------------------------------------------------------------------------
@@ -331,10 +331,10 @@ def replace_all(inputstr, all_to_replace, replace_to):
     return inputstr
 
 
-def highlight(code):
+def highlight(code, force=False):
     """Attempt to highlight Coconut code for the terminal."""
     from coconut.terminal import logger  # hide to remove circular deps
-    if logger.enable_colors(sys.stdout) and logger.enable_colors(sys.stderr):
+    if force or logger.enable_colors(sys.stdout) and logger.enable_colors(sys.stderr):
         try:
             from coconut.highlighter import highlight_coconut_for_terminal
         except ImportError: