Add BytesDeserializer

Author: Balaji Veeramani

src/sagemaker/deserializers.py

@@ -39,3 +39,24 @@ def deserialize(self, data, content_type):
     @abc.abstractmethod
     def ACCEPT(self):
         """The content type that is expected from the inference endpoint."""
+
+
+class BytesDeserializer(BaseDeserializer):
+    """Deserialize a stream of bytes into a bytes object."""
+
+    ACCEPT = "application/octet-stream"
+
+    def deserialize(self, data, content_type):
+        """Read a stream of bytes returned from an inference endpoint.
+
+        Args:
+            data (object): A stream of bytes.
+            content_type (str): The MIME type of the data.
+
+        Returns:
+            bytes: The bytes object read from the stream.
+        """
+        try:
+            return data.read()
+        finally:
+            data.close()

src/sagemaker/predictor.py

@@ -623,32 +623,6 @@ def __call__(self, stream, content_type):
 csv_deserializer = _CsvDeserializer()
 
 
-class BytesDeserializer(object):
-    """Return the response as an undecoded array of bytes.
-
-    Args:
-        accept (str): The Accept header to send to the server (optional).
-    """
-
-    def __init__(self, accept=None):
-        """
-        Args:
-            accept:
-        """
-        self.accept = accept
-
-    def __call__(self, stream, content_type):
-        """
-        Args:
-            stream:
-            content_type:
-        """
-        try:
-            return stream.read()
-        finally:
-            stream.close()
-
-
 class StringDeserializer(object):
     """Return the response as a decoded string.
 

tests/unit/sagemaker/cli/compatibility/v2/modifiers/test_predictors.py

@@ -117,9 +117,11 @@ def test_import_from_node_should_be_modified_random_import():
 def test_import_from_modify_node():
     modifier = predictors.PredictorImportFromRenamer()
 
-    node = ast_import("from sagemaker.predictor import BytesDeserializer, RealTimePredictor")
+    node = ast_import(
+        "from sagemaker.predictor import ClassThatHasntBeenRenamed, RealTimePredictor"
+    )
     modifier.modify_node(node)
-    expected_result = "from sagemaker.predictor import BytesDeserializer, Predictor"
+    expected_result = "from sagemaker.predictor import ClassThatHasntBeenRenamed, Predictor"
     assert expected_result == pasta.dump(node)
 
     node = ast_import("from sagemaker.predictor import RealTimePredictor as RTP")

tests/unit/test_predictor.py

@@ -20,13 +20,13 @@
 import pytest
 from mock import Mock, call, patch
 
+from sagemaker.deserializers import BytesDeserializer
 from sagemaker.predictor import Predictor
 from sagemaker.predictor import (
     json_serializer,
     json_deserializer,
     csv_serializer,
     csv_deserializer,
-    BytesDeserializer,
     StringDeserializer,
     StreamDeserializer,
     numpy_deserializer,
@@ -185,7 +185,7 @@ def test_json_deserializer_invalid_data():
 
 
 def test_bytes_deserializer():
-    result = BytesDeserializer()(io.BytesIO(b"[1, 2, 3]"), "application/json")
+    result = BytesDeserializer().deserialize(io.BytesIO(b"[1, 2, 3]"), "application/json")
 
     assert result == b"[1, 2, 3]"