From 7b81d31090b5c0282dc1ea0562fd9cbb7afb0d87 Mon Sep 17 00:00:00 2001
From: Bruno Oliveira <nicoddemus@gmail.com>
Date: Sun, 2 Feb 2020 09:23:41 -0300
Subject: [PATCH] Add "TerminalWriter.write_source" to write highlighted source
 code

---
 CHANGELOG                          |  6 ++++++
 py/_io/terminalwriter.py           | 13 +++++++++++++
 testing/io_/test_terminalwriter.py | 13 +++++++++++++
 tox.ini                            |  1 +
 4 files changed, 33 insertions(+)

diff --git a/CHANGELOG b/CHANGELOG
index 03a3d4af..97a1fd7e 100644
--- a/CHANGELOG
+++ b/CHANGELOG
@@ -1,3 +1,9 @@
+1.9.0 (UNRELEASED)
+==================
+
+- New ``TerminalWriter.write_source`` method that writes highlighted code to the terminal if
+  ``Pygments`` is installed.
+
 1.8.1 (2019-12-27)
 ==================
 
diff --git a/py/_io/terminalwriter.py b/py/_io/terminalwriter.py
index be559867..91e45be8 100644
--- a/py/_io/terminalwriter.py
+++ b/py/_io/terminalwriter.py
@@ -255,6 +255,19 @@ def write(self, msg, **kw):
                 markupmsg = msg
             write_out(self._file, markupmsg)
 
+    def write_source(self, source):
+        """Writes the given source code highlighted"""
+        if self.hasmarkup:
+            try:
+                from pygments.formatters.terminal import TerminalFormatter
+                from pygments.lexers.python import PythonLexer
+                from pygments import highlight
+            except ImportError:
+                pass
+            else:
+                source = highlight(source, PythonLexer(), TerminalFormatter())
+        self.write(source)
+
     def _update_chars_on_current_line(self, text_or_bytes):
         newline = b'\n' if isinstance(text_or_bytes, bytes) else '\n'
         current_line = text_or_bytes.rsplit(newline, 1)[-1]
diff --git a/testing/io_/test_terminalwriter.py b/testing/io_/test_terminalwriter.py
index 1eef7f7d..9fe605ab 100644
--- a/testing/io_/test_terminalwriter.py
+++ b/testing/io_/test_terminalwriter.py
@@ -303,3 +303,16 @@ def test_should_do_markup_PY_COLORS_eq_0(monkeypatch):
     tw.line("hello", bold=True)
     s = f.getvalue()
     assert s == "hello\n"
+
+
+@pytest.mark.parametrize('atty', [True, False])
+def test_write_source(atty):
+    f = py.io.TextIO()
+    f.isatty = lambda: atty
+    tw = py.io.TerminalWriter(file=f)
+    tw.write_source("x = 1")
+    s = f.getvalue()
+    if atty:
+        assert s == 'x = \x1b[34m1\x1b[39;49;00m\n'
+    else:
+        assert s == "x = 1"
diff --git a/tox.ini b/tox.ini
index 379249d3..f4261864 100644
--- a/tox.ini
+++ b/tox.ini
@@ -8,6 +8,7 @@ commands=
   py.test --confcutdir=. --junitxml={envlogdir}/junit-{envname}.xml []
 deps=
   attrs
+  pygments
   pytest29: pytest~=2.9.0
   pytest30: pytest~=3.0.0
   pytest31: pytest~=3.1.0