Preserve error type

Author: chriso

src/dispatch/proto.py

@@ -2,6 +2,7 @@
 
 import pickle
 from dataclasses import dataclass
+from types import TracebackType
 from typing import Any
 
 import google.protobuf.any_pb2
@@ -280,7 +281,13 @@ class Error:
     Output.
     """
 
-    def __init__(self, status: Status, type: str | None, message: str | None):
+    def __init__(
+        self,
+        status: Status,
+        type: str | None,
+        message: str | None,
+        value: Exception | None = None,
+    ):
         """Create a new Error.
 
         Args:
@@ -300,6 +307,7 @@ def __init__(self, status: Status, type: str | None, message: str | None):
         self.type = type
         self.message = message
         self.status = status
+        self.value = value
 
     @classmethod
     def from_exception(cls, ex: Exception, status: Status | None = None) -> Error:
@@ -313,18 +321,29 @@ def from_exception(cls, ex: Exception, status: Status | None = None) -> Error:
         if status is None:
             status = status_for_error(ex)
 
-        return Error(status, ex.__class__.__qualname__, str(ex))
+        return Error(status, ex.__class__.__qualname__, str(ex), ex)
 
     def to_exception(self) -> Exception:
-        # TODO: use correct error type
-        return RuntimeError(self.message)
+        if self.value is not None:
+            return self.value
+
+        g = globals()
+        try:
+            cls = g[self.type]
+            assert issubclass(cls, Exception)
+        except (KeyError, AssertionError):
+            return RuntimeError(self.message)
+        else:
+            return cls(self.message)
 
     @classmethod
     def _from_proto(cls, proto: error_pb.Error) -> Error:
-        return cls(Status.UNSPECIFIED, proto.type, proto.message)
+        value = pickle.loads(proto.value) if proto.value else None
+        return cls(Status.UNSPECIFIED, proto.type, proto.message, value)
 
     def _as_proto(self) -> error_pb.Error:
-        return error_pb.Error(type=self.type, message=self.message)
+        value = pickle.dumps(self.value) if self.value else None
+        return error_pb.Error(type=self.type, message=self.message, value=value)
 
 
 def _any_unpickle(any: google.protobuf.any_pb2.Any) -> Any:

tests/dispatch/test_scheduler.py

@@ -31,6 +31,11 @@ async def call_sequentially(*functions):
     return results
 
 
+@durable
+async def raises_error():
+    raise ValueError("oops")
+
+
 class TestOneShotScheduler(unittest.TestCase):
     def test_main_return(self):
         @durable
@@ -43,10 +48,10 @@ async def main():
     def test_main_raise(self):
         @durable
         async def main():
-            raise RuntimeError("oops")
+            raise ValueError("oops")
 
         output = self.start(main)
-        self.assert_exit_result_error(output, RuntimeError, "oops")
+        self.assert_exit_result_error(output, ValueError, "oops")
 
     def test_main_args(self):
         @durable
@@ -215,6 +220,14 @@ async def main():
 
         self.assertEqual(len(correlation_ids), 8)
 
+    def test_raise_indirect(self):
+        @durable
+        async def main():
+            return await gather(call_one("a"), raises_error())
+
+        output = self.start(main)
+        self.assert_exit_result_error(output, ValueError, "oops")
+
     def start(self, main: Callable, *args: Any, **kwargs: Any) -> Output:
         input = Input.from_input_arguments(main.__qualname__, *args, **kwargs)
         return OneShotScheduler(main).run(input)
@@ -258,6 +271,7 @@ def assert_exit_result_error(
         self.assertEqual(error.__class__, expect)
         if message is not None:
             self.assertEqual(str(error), message)
+        return error
 
     def assert_poll(self, output: Output) -> poll_pb.Poll:
         response = output._message