Repository: fal-ai/isolate (language: python)

Issue:
feat: consolidate serialization/deserialization for proper error handling
Stack from [ghstack](https://github.com/ezyang/ghstack) (oldest at bottom):
* #30
* #29
* #28
* #27
* #25
* #23
* #22
* #21
* __->__ #20
* #19




Discussion / hints:

Fix patch:
diff --git a/src/isolate/backends/connections/common.py b/src/isolate/backends/connections/common.py
new file mode 100644
index 0000000..9f68cbb
--- /dev/null
+++ b/src/isolate/backends/connections/common.py
@@ -0,0 +1,84 @@
+from __future__ import annotations
+
+import importlib
+from contextlib import contextmanager
+from typing import TYPE_CHECKING, Any, Iterator, cast
+
+if TYPE_CHECKING:
+    from typing import Protocol
+
+    class SerializationBackend(Protocol):
+        def loads(self, data: bytes) -> Any:
+            ...
+
+        def dumps(self, obj: Any) -> bytes:
+            ...
+
+
+class SerializationError(Exception):
+    """An error that happened during the serialization process."""
+
+
+@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."""
+
+    try:
+        yield
+    except BaseException as exception:
+        raise SerializationError("Error while " + message) from exception
+
+
+def as_serialization_backend(backend: Any) -> SerializationBackend:
+    """Ensures that the given backend has loads/dumps methods, and returns
+    it as is (also convinces type checkers that the given object satisfies
+    the serialization protocol)."""
+
+    if not hasattr(backend, "loads") or not hasattr(backend, "dumps"):
+        raise TypeError(
+            f"The given serialization backend ({backend.__name__}) does "
+            "not have one of the required methods (loads/dumps)."
+        )
+
+    return cast("SerializationBackend", backend)
+
+
+def load_serialized_object(
+    serialization_method: str,
+    raw_object: bytes,
+    *,
+    was_it_raised: bool = False,
+) -> Any:
+    """Load the given serialized object using the given serialization method. If
+    anything fails, then a SerializationError will be raised. If the was_it_raised
+    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})"):
+        serialization_backend = as_serialization_backend(
+            importlib.import_module(serialization_method)
+        )
+
+    with _step("deserializing the given object"):
+        result = serialization_backend.loads(raw_object)
+
+    if was_it_raised:
+        raise result
+    else:
+        return result
+
+
+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})"):
+        serialization_backend = as_serialization_backend(
+            importlib.import_module(serialization_method)
+        )
+
+    with _step("serializing the given object"):
+        return serialization_backend.dumps(object)


Test patch:
diff --git a/tests/test_serialization.py b/tests/test_serialization.py
new file mode 100644
index 0000000..cca2fc8
--- /dev/null
+++ b/tests/test_serialization.py
@@ -0,0 +1,66 @@
+import pytest
+from functools import partial
+from isolate.backends.connections.common import (
+    SerializationError,
+    as_serialization_backend,
+    load_serialized_object,
+    serialize_object,
+)
+
+
+@pytest.mark.parametrize(
+    "method",
+    [
+        "pickle",
+        "dill",
+        "cloudpickle",
+    ],
+)
+def test_serialize_object(method):
+    func = partial(eval, "2 + 2")
+    serialized = serialize_object(method, func)
+    deserialized = load_serialized_object(method, serialized)
+    assert deserialized() == 4
+
+
+def test_deserialize_exception():
+    serialized = serialize_object("pickle", ValueError("some error"))
+    regular_obj = load_serialized_object("pickle", serialized)
+    assert isinstance(regular_obj, ValueError)
+    assert regular_obj.args == ("some error",)
+
+
+def test_deserialize_raised_exception():
+    serialized = serialize_object("pickle", ValueError("some error"))
+    with pytest.raises(ValueError) as exc_info:
+        load_serialized_object("pickle", serialized, was_it_raised=True)
+    assert exc_info.value.args == ("some error",)
+
+
+def error_while_serializing():
+    anon = lambda: 2 + 2  # anonymous functions are not
+    # serializable by pickle
+    with pytest.raises(SerializationError) as exc_info:
+        serialize_object("pickle", anon)
+
+    assert exc_info.match("Error while serializing the given object")
+
+    dill_serialized_lambda = serialize_object("dill", anon)
+
+    with pytest.raises(SerializationError) as exc_info:
+        load_serialized_object("pickle", dill_serialized_lambda)
+
+    assert exc_info.match("Error while deserializing the given object")
+
+
+
+def error_while_loading_backend():
+    with pytest.raises(SerializationError) as exc_info:
+        serialize_object("$$$", 1)
+
+    assert exc_info.match("Error while preparing the serialization backend")
+
+    with pytest.raises(SerializationError) as exc_info:
+        load_serialized_object("$$$", b"1")
+
+    assert exc_info.match("Error while preparing the serialization backend")