diagnostic.Engine: allow adding multiple notes at once.

m-labs · Jul 7, 2016 · c744600 · c744600
1 parent 5cb107b
commit c744600
Showing 2 changed files with 16 additions and 6 deletions.
diff --git a/pythonparser/diagnostic.py b/pythonparser/diagnostic.py
@@ -165,14 +165,14 @@ def process(self, diagnostic):
             raise Error(diagnostic)
 
     @contextmanager
-    def context(self, note):
+    def context(self, *notes):
         """
         A context manager that appends ``note`` to every diagnostic processed by
         this engine.
         """
-        self._appended_notes.append(note)
+        self._appended_notes += notes
         yield
-        self._appended_notes.pop()
+        del self._appended_notes[-len(notes):]
 
     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
@@ -37,16 +37,26 @@ def render_diagnostic(diag):
         self.engine.render_diagnostic = render_diagnostic
 
     def test_context(self):
-        note = diagnostic.Diagnostic(
+        note1 = diagnostic.Diagnostic(
             "note", "broken here", {},
             source.Range(self.buffer, 0, 0))
+        note2 = diagnostic.Diagnostic(
+            "note", "also broken there", {},
+            source.Range(self.buffer, 0, 0))
+
+        with self.engine.context(note1):
+            with self.engine.context(note2):
+                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, [note1, note2])
 
-        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])
+            self.assertEqual(self.last_diagnostic.notes, [note1])
 
         diag = diagnostic.Diagnostic(
             "error", "{x} doesn't work", {"x": "everything"},