m-labs · Jul 7, 2016 · Jul 7, 2016
Showing with 47 additions and 1 deletion.

+19 −1 pythonparser/diagnostic.py

+28 −0 pythonparser/test/test_diagnostic.py
diff --git a/pythonparser/diagnostic.py b/pythonparser/diagnostic.py
@@ -5,6 +5,7 @@
 
 from __future__ import absolute_import, division, print_function, unicode_literals
 from functools import reduce
+from contextlib import contextmanager
 import sys, re
 
 class Diagnostic:
@@ -38,10 +39,15 @@ class Diagnostic:
     """
 
     def __init__(self, level, reason, arguments, location,
-                 highlights=[], notes=[]):
+                 highlights=None, notes=None):
         if level not in self.LEVELS:
             raise ValueError("level must be one of Diagnostic.LEVELS")
 
+        if highlights is None:
+            highlights = []
+        if notes is None:
+            notes = []
+
         if len(set(map(lambda x: x.source_buffer,
                        [location] + highlights))) > 1:
             raise ValueError("location and highlights must refer to the same source buffer")
@@ -145,16 +151,28 @@ class Engine:
     """
     def __init__(self, all_errors_are_fatal=False):
         self.all_errors_are_fatal = all_errors_are_fatal
+        self._appended_notes = []
 
     def process(self, diagnostic):
         """
         The default implementation of :meth:`process` renders non-fatal
         diagnostics to ``sys.stderr``, and raises fatal ones as a :class:`Error`.
         """
+        diagnostic.notes += self._appended_notes
         self.render_diagnostic(diagnostic)
         if diagnostic.level == "fatal" or \
                 (self.all_errors_are_fatal and diagnostic.level == "error"):
             raise Error(diagnostic)
 
+    @contextmanager
+    def context(self, note):
+        """
+        A context manager that appends ``note`` to every diagnostic processed by
+        this engine.
+        """
+        self._appended_notes.append(note)
+        yield
+        self._appended_notes.pop()
+
     def render_diagnostic(self, diagnostic):
         sys.stderr.write("\n".join(diagnostic.render()) + "\n")
diff --git a/pythonparser/test/test_diagnostic.py b/pythonparser/test/test_diagnostic.py
@@ -25,3 +25,31 @@ def test_render(self):
              "x + (1 + 'a')",
              "     ~ ^ ~~~ "],
             diag.render())
+
+class DiagnosticEngineTestCase(unittest.TestCase):
+
+    def setUp(self):
+        self.buffer = source.Buffer("x + (1 + 'a')\n")
+        self.last_diagnostic = None
+        self.engine = diagnostic.Engine()
+        def render_diagnostic(diag):
+            self.last_diagnostic = diag
+        self.engine.render_diagnostic = render_diagnostic
+
+    def test_context(self):
+        note = diagnostic.Diagnostic(
+            "note", "broken here", {},
+            source.Range(self.buffer, 0, 0))
+
+        with self.engine.context(note):
+            diag = diagnostic.Diagnostic(
+                "error", "{x} doesn't work", {"x": "everything"},
+                source.Range(self.buffer, 0, 0))
+            self.engine.process(diag)
+            self.assertEqual(self.last_diagnostic.notes, [note])
+
+        diag = diagnostic.Diagnostic(
+            "error", "{x} doesn't work", {"x": "everything"},
+            source.Range(self.buffer, 0, 0))
+        self.engine.process(diag)
+        self.assertEqual(self.last_diagnostic.notes, [])