Add include_traceback option to SourceLine. (#133)

Author: tetron

schema_salad/sourceline.py

@@ -3,6 +3,7 @@
 from ruamel.yaml.comments import CommentedBase, CommentedMap, CommentedSeq
 import re
 import os
+import traceback
 
 from typing import (Any, AnyStr, Callable, cast, Dict, List, Iterable, Tuple,
                     TypeVar, Union, Text)
@@ -134,22 +135,26 @@ def cmap(d, lc=None, fn=None):  # type: (Union[int, float, str, Text, Dict, List
         return d
 
 class SourceLine(object):
-    def __init__(self, item, key=None, raise_type=six.text_type):  # type: (Any, Any, Callable) -> None
+    def __init__(self, item, key=None, raise_type=six.text_type, include_traceback=False):  # type: (Any, Any, Callable, bool) -> None
         self.item = item
         self.key = key
         self.raise_type = raise_type
+        self.include_traceback = include_traceback
 
     def __enter__(self):  # type: () -> SourceLine
         return self
 
     def __exit__(self,
                  exc_type,   # type: Any
                  exc_value,  # type: Any
-                 traceback   # type: Any
-                 ):  # -> Any
+                 tb   # type: Any
+                 ):   # -> Any
         if not exc_value:
             return
-        raise self.makeError(six.text_type(exc_value))
+        if self.include_traceback:
+            raise self.makeError("\n".join(traceback.format_exception(exc_type, exc_value, tb)))
+        else:
+            raise self.makeError(six.text_type(exc_value))
 
     def makeLead(self):  # type: () -> Text
         if self.key is None or self.item.lc.data is None or self.key not in self.item.lc.data:

schema_salad/tests/test_examples.py

@@ -10,7 +10,7 @@
 import ruamel.yaml
 import json
 import os
-from schema_salad.sourceline import cmap
+from schema_salad.sourceline import cmap, SourceLine
 
 try:
     from ruamel.yaml import CSafeLoader as SafeLoader
@@ -376,5 +376,30 @@ def test_file_uri(self):
                           schema_salad.ref_resolver.uri_file_path("file:///foo/bar%20baz/quux#zing%20zong"))
 
 
+class SourceLineTest(unittest.TestCase):
+    def test_sourceline(self):
+        ldr = schema_salad.ref_resolver.Loader({"id": "@id"})
+        b, _ = ldr.resolve_ref(get_data("tests/frag.yml"))
+
+        class TestExp(Exception):
+            pass
+
+        try:
+            with SourceLine(b, 1, TestExp, False):
+                raise Exception("Whoops")
+        except TestExp as e:
+            self.assertTrue(str(e).endswith("frag.yml:3:3: Whoops"))
+        except:
+            self.assertFail()
+
+        try:
+            with SourceLine(b, 1, TestExp, True):
+                raise Exception("Whoops")
+        except TestExp as e:
+            self.assertTrue(str(e).splitlines()[0].endswith("frag.yml:3:3: Traceback (most recent call last):"))
+        except:
+            self.assertFail()
+
+
 if __name__ == '__main__':
     unittest.main()