Merge branch 'zwei' into add-string-deserializer

Author: bveeramani

src/sagemaker/deserializers.py

@@ -68,3 +68,24 @@ def deserialize(self, data, content_type):
             return data.read().decode(self.encoding)
         finally:
             data.close()
+
+
+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 StreamDeserializer(object):
     """Returns the tuple of the response stream and the content-type of the response.
        It is the receivers responsibility to close the stream when they're done

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/sagemaker/test_deserializers.py

@@ -14,7 +14,7 @@
 
 import io
 
-from sagemaker.deserializers import StringDeserializer
+from sagemaker.deserializers import StringDeserializer, BytesDeserializer
 
 
 def test_string_deserializer():
@@ -23,3 +23,11 @@ def test_string_deserializer():
     result = deserializer.deserialize(io.BytesIO(b"[1, 2, 3]"), "application/json")
 
     assert result == "[1, 2, 3]"
+
+    
+def test_bytes_deserializer():
+    deserializer = BytesDeserializer()
+
+    result = deserializer.deserialize(io.BytesIO(b"[1, 2, 3]"), "application/json")
+
+    assert result == b"[1, 2, 3]"

tests/unit/test_predictor.py

@@ -26,7 +26,6 @@
     json_deserializer,
     csv_serializer,
     csv_deserializer,
-    BytesDeserializer,
     StreamDeserializer,
     numpy_deserializer,
     npy_serializer,
@@ -181,13 +180,7 @@ def test_json_deserializer_invalid_data():
     with pytest.raises(ValueError) as error:
         json_deserializer(io.BytesIO(b"[[1]"), "application/json")
     assert "column" in str(error)
-
-
-def test_bytes_deserializer():
-    result = BytesDeserializer()(io.BytesIO(b"[1, 2, 3]"), "application/json")
-
-    assert result == b"[1, 2, 3]"
-
+    
 
 def test_stream_deserializer():
     stream, content_type = StreamDeserializer()(io.BytesIO(b"[1, 2, 3]"), "application/json")