extend the Encoder/Decoder state API to be type-aware

thrau · thrau · commit 48c788c8ea44 · 2025-08-27T12:12:11.000+02:00
diff --git a/localstack-core/localstack/state/core.py b/localstack-core/localstack/state/core.py
@@ -97,42 +97,48 @@ def __str__(self) -> str:
 
 
 class Encoder:
-    def encodes(self, obj: Any) -> bytes:
+    def encodes(self, obj: Any, py_type: type = None) -> bytes:
         """
         Encode an object into bytes.
 
         :param obj: the object to encode
+        :param py_type: the type of the object. needed by some encoders that don't have implicit type knowledge.
         :return: the encoded object
         """
         b = io.BytesIO()
         self.encode(obj, b)
         return b.getvalue()
 
-    def encode(self, obj: Any, file: IO[bytes]):
+    def encode(self, obj: Any, file: IO[bytes], py_type: type = None):
         """
         Encode an object into bytes.
 
         :param obj: the object to encode
+        :param py_type: the type of the object. needed by some encoders that don't have implicit type knowledge.
         :param file: the file to write the encoded data into
         """
         raise NotImplementedError
 
 
 class Decoder:
-    def decodes(self, data: bytes) -> Any:
+    def decodes(self, data: bytes, py_type: type = None) -> Any:
         """
         Decode a previously encoded object.
 
         :param data: the encoded object to decode
+        :param py_type: the type that is expected as return type. Needed by some decoders that don't have implicit
+        type knowledge.
         :return: the decoded object
         """
-        return self.decode(io.BytesIO(data))
+        return self.decode(io.BytesIO(data), py_type)
 
-    def decode(self, file: IO[bytes]) -> Any:
+    def decode(self, file: IO[bytes], py_type: type = None) -> Any:
         """
         Decode a previously encoded object.
 
         :param file: the io object containing the object to decode
+        :param py_type: the type that is expected as return type. Needed by some decoders that don't have implicit
+        type knowledge.
         :return: the decoded object
         """
         raise NotImplementedError
diff --git a/localstack-core/localstack/state/pickle.py b/localstack-core/localstack/state/pickle.py
@@ -240,7 +240,7 @@ class PickleEncoder(Encoder):
     def __init__(self, pickler_class: type[dill.Pickler] = None):
         self.pickler_class = pickler_class or Pickler
 
-    def encode(self, obj: Any, file: BinaryIO):
+    def encode(self, obj: Any, file: BinaryIO, py_type: type = None) -> Any:
         return self.pickler_class(file).dump(obj)
 
 
@@ -255,7 +255,7 @@ class PickleDecoder(Decoder):
     def __init__(self, unpickler_class: type[dill.Unpickler] = None):
         self.unpickler_class = unpickler_class or dill.Unpickler
 
-    def decode(self, file: BinaryIO) -> Any:
+    def decode(self, file: BinaryIO, py_type=None) -> Any:
         return self.unpickler_class(file).load()
 
 
@@ -267,7 +267,7 @@ class JsonEncoder(Encoder):
     def __init__(self, pickler_class: type[jsonpickle.Pickler] = None):
         self.pickler_class = pickler_class or jsonpickle.Pickler()
 
-    def encode(self, obj: Any, file: IO[bytes]):
+    def encode(self, obj: Any, file: IO[bytes], py_type: type = None):
         json_str = jsonpickle.encode(obj, context=self.pickler_class)
         file.write(json_str.encode("utf-8"))
 
@@ -282,7 +282,7 @@ class JsonDecoder(Decoder):
     def __init__(self, unpickler_class: type[jsonpickle.Unpickler] = None):
         self.unpickler_class = unpickler_class or jsonpickle.Unpickler()
 
-    def decode(self, file: IO[bytes]) -> Any:
+    def decode(self, file: IO[bytes], py_type: type = None) -> Any:
         json_str = file.read().decode("utf-8")
         return jsonpickle.decode(json_str, context=self.unpickler_class)
 
