From 357e05e8550c8537cf0842e90cc6af9764475251 Mon Sep 17 00:00:00 2001
From: Ruslan Kuprieiev <kupruser@gmail.com>
Date: Tue, 14 May 2024 16:57:07 +0300
Subject: [PATCH] chore(isolate): separate serialization and deserialization
 exceptions (#129)

---
 src/isolate/connections/common.py | 43 +++++++++++++++++++------------
 1 file changed, 26 insertions(+), 17 deletions(-)

diff --git a/src/isolate/connections/common.py b/src/isolate/connections/common.py
index 3d472f6..954ded2 100644
--- a/src/isolate/connections/common.py
+++ b/src/isolate/connections/common.py
@@ -2,9 +2,8 @@
 
 import importlib
 import os
-from contextlib import contextmanager
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Any, Iterator, cast
+from typing import TYPE_CHECKING, Any, cast
 
 from tblib import Traceback, TracebackParseError
 
@@ -23,23 +22,18 @@ def dumps(self, obj: Any) -> bytes: ...
 
 
 @dataclass
-class SerializationError(IsolateException):
+class BaseSerializationError(IsolateException):
     """An error that happened during the serialization process."""
 
     message: str
 
 
-@contextmanager
-def _step(message: str) -> Iterator[None]:
-    """A context manager to capture every expression
-    underneath it and if any of them fails for any reason
-    then it will raise a SerializationError with the
-    given message."""
+class SerializationError(BaseSerializationError):
+    pass
 
-    try:
-        yield
-    except BaseException as exception:
-        raise SerializationError("Error while " + message) from exception
+
+class DeserializationError(BaseSerializationError):
+    pass
 
 
 def as_serialization_method(backend: Any) -> SerializationBackend:
@@ -68,13 +62,22 @@ def load_serialized_object(
     flag is set to true, then the given object will be raised as an exception (instead
     of being returned)."""
 
-    with _step(f"preparing the serialization backend ({serialization_method})"):
+    try:
         serialization_backend = as_serialization_method(
             importlib.import_module(serialization_method)
         )
+    except BaseException as exc:
+        raise DeserializationError(
+            "Error while preparing the serialization backend "
+            f"({serialization_method})"
+        ) from exc
 
-    with _step("deserializing the given object"):
+    try:
         result = serialization_backend.loads(raw_object)
+    except BaseException as exc:
+        raise DeserializationError(
+            "Error while deserializing the given object"
+        ) from exc
 
     if was_it_raised:
         raise prepare_exc(result, stringized_traceback=stringized_traceback)
@@ -86,13 +89,19 @@ def serialize_object(serialization_method: str, object: Any) -> bytes:
     """Serialize the given object using the given serialization method. If
     anything fails, then a SerializationError will be raised."""
 
-    with _step(f"preparing the serialization backend ({serialization_method})"):
+    try:
         serialization_backend = as_serialization_method(
             importlib.import_module(serialization_method)
         )
+    except BaseException as exc:
+        raise SerializationError(
+            f"Error while preparing the serialization backend ({serialization_method})"
+        ) from exc
 
-    with _step("serializing the given object"):
+    try:
         return serialization_backend.dumps(object)
+    except BaseException as exc:
+        raise SerializationError("Error while serializing the given object") from exc
 
 
 def is_agent() -> bool: