Fix issues when rendering exception traces

- Fixes the highlighting of empty sources
- Fixes the rendering of docstrings

Author: sdispater

cleo/ui/exception_trace.py

@@ -104,11 +104,17 @@ def readline():
 
             if token_type == tokenize.ENDMARKER:
                 # End of source
+                if current_type is None:
+                    current_type = self.TOKEN_DEFAULT
+
                 line += "<{}>{}</>".format(self._theme[current_type], buffer)
                 lines.append(line)
                 break
 
             if lineno > current_line:
+                if current_type is None:
+                    current_type = self.TOKEN_DEFAULT
+
                 diff = lineno - current_line
                 if diff > 1:
                     lines += [""] * (diff - 1)
@@ -154,8 +160,9 @@ def readline():
 
             if lineno < end[0]:
                 # The token spans multiple lines
-                lines.append(line)
                 token_lines = token_string.split("\n")
+                line += "<{}>{}</>".format(self._theme[current_type], token_lines[0])
+                lines.append(line)
                 for token_line in token_lines[1:-1]:
                     lines.append(
                         "<{}>{}</>".format(self._theme[current_type], token_line)

tests/ui/test_exception_trace.py

@@ -443,3 +443,37 @@ def call():
     )
 
     assert expected == io.fetch_output()
+
+
+def test_empty_source_file_do_not_break_highlighter():
+    from cleo.ui.exception_trace import Highlighter
+
+    highlighter = Highlighter()
+    highlighter.highlighted_lines("")
+
+
+def test_doctrings_are_corrrectly_rendered():
+    from cleo.formatters.formatter import Formatter
+    from cleo.ui.exception_trace import Highlighter
+
+    source = '''
+def test():
+    """
+    Doctring
+    """
+    ...
+'''
+
+    formatter = Formatter()
+    highlighter = Highlighter()
+    lines = highlighter.highlighted_lines(source)
+
+    assert [formatter.format(l) for l in lines] == [
+        "",
+        "def test():",
+        '    """',
+        "    Doctring",
+        '    """',
+        "    ...",
+        "",
+    ]