ansi improvements and tty detection

Author: samuelcolvin

devtools/ansi.py

@@ -1,13 +1,21 @@
+import sys
 from enum import IntEnum
 
 _ansi_template = '\033[{}m'
 
 
+def isatty(stream):
+    try:
+        return stream.isatty()
+    except Exception:
+        return False
+
+
 class Style(IntEnum):
     """
     Heavily borrowed from https://github.com/pallets/click/blob/6.7/click/termui.py
 
-    Italic added and generally modernised improved.
+    Italic added, multiple ansi codes condensed into one block and generally modernised.
     """
     reset = 0
 
@@ -85,18 +93,23 @@ def __call__(self, text: str, *styles, reset: bool=True):
                       string which means that styles do not carry over.  This
                       can be disabled to compose styles.
         """
-        parts = []
+        codes = []
         for s in styles:
             if not isinstance(s, self.__class__):
                 try:
                     s = self.styles[s]
                 except KeyError:
                     raise ValueError('invalid style "{}"'.format(s))
-            parts.append(_ansi_template.format(s))
-        parts.append(text)
+            codes.append(str(s.value))
+
+        if codes:
+            r = _ansi_template.format(';'.join(codes)) + text
+        else:
+            r = text
+
         if reset:
-            parts.append(_ansi_template.format(self.reset))
-        return ''.join(parts)
+            r += _ansi_template.format(self.reset)
+        return r
 
     @property
     def styles(self):
@@ -118,5 +131,7 @@ def __str__(self):
 style = Style(-1)
 
 
-def sprint(text, *styles, reset=True, flush=True, **print_kwargs):
-    print(style(text, *styles, reset=reset), flush=flush, **print_kwargs)
+def sprint(text, *styles, reset=True, flush=True, file=None, **print_kwargs):
+    if isatty(file or sys.stdout):
+        text = style(text, *styles, reset=reset)
+    print(text, flush=flush, file=file, **print_kwargs)

devtools/debug.py

@@ -56,12 +56,8 @@ def __init__(self, *, filename: str, lineno: int, frame: str, arguments: List[De
         self.arguments = arguments
 
     def __str__(self) -> str:
-        template = '{s.filename}:{s.lineno} {s.frame}'
-        # turns out single line output is ugly
-        # if len(self.arguments) == 1:
-        #     return (template + ': {a}').format(s=self, a=self.arguments[0])
-        # else:
-        return (template + '\n  {a}').format(s=self, a='\n  '.join(str(a) for a in self.arguments))
+        template = '{s.filename}:{s.lineno} {s.frame}\n  {a}'
+        return template.format(s=self, a='\n  '.join(str(a) for a in self.arguments))
 
     def __repr__(self) -> str:
         arguments = ' '.join(str(a) for a in self.arguments)
@@ -78,14 +74,20 @@ class Debug:
 
     def __init__(self, *,
                  warnings: Optional[bool]=None,
+                 colours: Optional[bool]=None,
                  frame_context_length: int=50):
-        if warnings is None:
-            self._warnings = env_true('PY_DEVTOOLS_WARNINGS', 'TRUE')
-        else:
-            self._warnings = warnings
+        self._warnings = self._env_bool(warnings, 'PY_DEVTOOLS_WARNINGS')
+        self._colours = self._env_bool(colours, 'PY_DEVTOOLS_COLOURS')
         # 50 lines should be enough to make sure we always get the entire function definition
         self._frame_context_length = frame_context_length
 
+    @classmethod
+    def _env_bool(cls, value, env_name, env_default='TRUE'):
+        if value is None:
+            return env_true(env_name, env_default)
+        else:
+            return value
+
     def __call__(self, *args, **kwargs):
         print(self._process(args, kwargs, r'debug *\('), flush=True)
 

tests/test_ansi.py

@@ -1,21 +1,33 @@
+import io
+
 import pytest
 
 from devtools.ansi import sprint, style
 
 
 def test_colorize():
     v = style('hello', style.red)
-    assert v == '\x1b[31mhello\x1b[0m', repr(v)
+    assert v == '\x1b[31mhello\x1b[0m'
 
 
 def test_no_reset():
     v = style('hello', style.bold, reset=False)
-    assert v == '\x1b[1mhello', repr(v)
+    assert v == '\x1b[1mhello'
+
+
+def test_combine_styles():
+    v = style('hello', style.red, style.bold)
+    assert v == '\x1b[31;1mhello\x1b[0m'
+
+
+def test_no_styles():
+    v = style('hello')
+    assert v == 'hello\x1b[0m'
 
 
 def test_style_str():
     v = style('hello', 'red')
-    assert v == '\x1b[31mhello\x1b[0m', repr(v)
+    assert v == '\x1b[31mhello\x1b[0m'
 
 
 def test_invalid_style_str():
@@ -24,11 +36,33 @@ def test_invalid_style_str():
     assert exc_info.value.args[0] == 'invalid style "mauve"'
 
 
-def test_print(capsys):
-    sprint('hello', style.green)
-    stdout, stderr = capsys.readouterr()
-    assert stdout == '\x1b[32mhello\x1b[0m\n', repr(stdout)
-    assert stderr == ''
+def test_print_not_tty():
+    stream = io.StringIO()
+    sprint('hello', style.green, file=stream)
+    out = stream.getvalue()
+    assert out == 'hello\n'
+
+
+def test_print_is_tty():
+    class TTYStream(io.StringIO):
+        def isatty(self):
+            return True
+
+    stream = TTYStream()
+    sprint('hello', style.green, file=stream)
+    out = stream.getvalue()
+    assert out == '\x1b[32mhello\x1b[0m\n', repr(out)
+
+
+def test_print_tty_error():
+    class TTYStream(io.StringIO):
+        def isatty(self):
+            raise RuntimeError('boom')
+
+    stream = TTYStream()
+    sprint('hello', style.green, file=stream)
+    out = stream.getvalue()
+    assert out == 'hello\n'
 
 
 def test_get_styles():