Merge 226827a into 8833026

mypy.ini

@@ -4,6 +4,8 @@ ignore_missing_imports = True
 ignore_missing_imports = True
 [mypy-psutil.*]
 ignore_missing_imports = True
+[mypy-pygments.*]
+ignore_missing_imports = True
 [mypy-testslide.testslide]
 disallow_untyped_calls = True
 disallow_untyped_defs = True

setup.py

@@ -22,6 +22,7 @@
     setup_requires=["setuptools>=38.6.0"],
     install_requires=[
         "psutil>=5.6.7",
+        "Pygments>=2.7.1",
         "typeguard>=2.9.1",
         'dataclasses==0.6; python_version < "3.7"',
     ],

tests/cli_unittest.py

@@ -271,21 +271,18 @@ def test_prints_exceptions_with_cause(self):
             tty_stdout=False,
             expected_return_code=1,
             expected_regex_in_stdout=(
-                (
-                    ".*1\) AssertionError: Third\n"
-                    ".*\n *"
-                    + '  File "tests/sample_tests.py", line [0-9]+, in test_failing\n *'
-                    + '    raise AssertionError\("Third"\) from cause\n *'
-                    + "\n *"
-                    + "  Caused by AssertionError: Second\n *"
-                    + '    File "tests/sample_tests.py", line [0-9]+, in test_failing\n *'
-                    + '      raise AssertionError\("Second"\) from cause\n *'
-                    + "\n *"
-                    + "    Caused by AssertionError: First\n *"
-                    + '      File "tests/sample_tests.py", line [0-9]+, in test_failing\n *'
-                    + '        raise AssertionError\("First"\)\n *'
-                    + ".*"
-                )
+                ".*\n"
+                "    1\) AssertionError: Third\n"
+                ".*\n"
+                '      File "tests/sample_tests.py", line [0-9]+, in test_failing\n'
+                '        raise AssertionError\("Third"\) from cause\n'
+                "      Caused by AssertionError: Second\n"
+                '        File "tests/sample_tests.py", line [0-9]+, in test_failing\n'
+                '          raise AssertionError\("Second"\) from cause\n'
+                "        Caused by AssertionError: First\n"
+                '          File "tests/sample_tests.py", line [0-9]+, in test_failing\n'
+                '            raise AssertionError\("First"\)\n *'
+                ".*"
             ),
         )
 

testslide/runner.py

@@ -15,6 +15,9 @@
 from typing import Any, Callable, Dict, List, Optional, Pattern, Union, cast
 
 import psutil
+import pygments
+import pygments.formatters
+import pygments.lexers
 
 from . import AggregatedExceptions, Context, Example, Skip, _ExampleRunner
 
@@ -161,9 +164,15 @@ def dsl_function(self, example: Example, code: Callable) -> None:
 
 
 class ColorFormatterMixin(BaseFormatter):
+    @property
+    def colored(self):
+        return sys.stdout.isatty() or self.force_color
+
     def _print_attrs(self, attrs: str, *values: Any, **kwargs: Any) -> None:
-        stream = kwargs.get("file", sys.stdout)
-        if stream.isatty() or self.force_color:
+        file = kwargs.get("file", None)
+        if file is not None:
+            raise ValueError()
+        if self.colored:
             print(
                 "\033[0m\033[{attrs}m{value}\033[0m".format(
                     attrs=attrs, value="".join([str(value) for value in values])
@@ -190,10 +199,26 @@ def print_cyan(self, *values: Any, **kwargs: Any) -> None:
 
 
 class FailurePrinterMixin(ColorFormatterMixin):
+    def _get_test_module_index(self, tb: traceback.StackSummary) -> int:
+        test_module_index = len(tb) - 1
+
+        test_module_paths = [
+            import_module(import_module_name).__file__
+            for import_module_name in self.import_module_names
+        ]
+
+        for index, value in enumerate(tb):
+            path = value[0]
+            if path in test_module_paths:
+                if index < test_module_index:
+                    test_module_index = index
+
+        return test_module_index
+
     def _print_stack_trace(self, exception: BaseException, cause_depth: int) -> None:
         indent = "  " * cause_depth
         if cause_depth:
-            self.print_red(f"\n    {indent}Caused by ", end="")
+            self.print_red(f"{indent}    Caused by ", end="")
 
         self.print_red(
             "{exception_class}: {message}".format(
@@ -204,37 +229,35 @@ def _print_stack_trace(self, exception: BaseException, cause_depth: int) -> None
 
         tb = traceback.extract_tb(exception.__traceback__)
 
-        test_module_index = len(tb) - 1
-
-        test_module_paths = [
-            import_module(import_module_name).__file__
-            for import_module_name in self.import_module_names
-        ]
-
-        for index, value in enumerate(tb):
-            path = value[0]
-            if path in test_module_paths:
-                if index < test_module_index:
-                    test_module_index = index
+        test_module_index = self._get_test_module_index(tb)
 
-        for index, value in enumerate(tb):
-            path, line, function_name, text = value
+        for index, (path, line, function_name, text) in enumerate(tb):
             if not self.show_testslide_stack_trace and index < test_module_index:
                 continue
             if self.trim_path_prefix:
                 split = path.split(self.trim_path_prefix)
                 if len(split) == 2 and not split[0]:
                     path = split[1]
-            self.print_cyan(
-                '{indent}      File "{path}", line {line}, in {function_name}\n'
-                "{indent}        {text}".format(
-                    indent=indent,
+            row_text = (
+                '  File "{path}", line {line}, in {function_name}\n'
+                "    {text}\n".format(
                     path=path,
                     line=line,
                     function_name=function_name,
                     text=text,
                 )
             )
+            if self.colored:
+                row_text = pygments.highlight(
+                    row_text,
+                    pygments.lexers.PythonTracebackLexer(),
+                    pygments.formatters.TerminalFormatter(),
+                )
+            row_text = "\n".join(
+                "{indent}    {line}".format(indent=indent, line=line)
+                for line in row_text.split("\n")[:-1]
+            )
+            print(row_text)
 
         if exception.__cause__:
             self._print_stack_trace(exception.__cause__, cause_depth=cause_depth + 1)