from pygments import highlight
from pygments.lexers import PythonTracebackLexer
from pygments.formatters import TerminalFormatter

from django.utils.unittest import TextTestRunner, TextTestResult

from discover_runner import DiscoverRunner


class HighlightedTextTestResult(TextTestResult):

    def _exc_info_to_string(self, err, test):
        code = super(HighlightedTextTestResult, self)._exc_info_to_string(err, test)
        return highlight(code, PythonTracebackLexer(), TerminalFormatter())


class HighlightedTextTestRunner(TextTestRunner):
    resultclass = HighlightedTextTestResult


class HighlightedDiscoverRunner(DiscoverRunner):

    def run_suite(self, suite, **kwargs):
        return HighlightedTextTestRunner(
            verbosity=self.verbosity, failfast=self.failfast).run(suite)