Add StringDeserializer

Author: Balaji Veeramani

src/sagemaker/deserializers.py

@@ -15,6 +15,8 @@
 
 import abc
 
+from sagemaker.utils import parse_mime_type
+
 
 class BaseDeserializer(abc.ABC):
     """Abstract base class for creation of new deserializers.
@@ -39,3 +41,37 @@ def deserialize(self, data, content_type):
     @abc.abstractmethod
     def ACCEPT(self):
         """The content type that is expected from the inference endpoint."""
+
+
+class StringDeserializer(object):
+    """Deserialize data from an inference endpoint into a decoded string."""
+
+    def __init__(self, encoding="UTF-8"):
+        """Initialize the default encoding.
+
+        Args:
+            encoding (str): The string encoding to use, if a charset is not
+                provided by the server (default: UTF-8).
+        """
+        self.encoding = encoding
+
+    def deserialize(self, data, content_type):
+        """Deserialize data from an inference endpoint into a decoded string.
+
+        Args:
+            data (object): A string or a byte stream.
+            content_type (str): The MIME type of the data.
+
+        Returns:
+            str: The data deserialized into a decoded string.
+        """
+        category, _, parameters = parse_mime_type(content_type)
+
+        if category == "text":
+            return data
+
+        try:
+            encoding = parameters.get("charset", self.encoding)
+            return data.read().decode(encoding)
+        finally:
+            data.close()

src/sagemaker/predictor.py

@@ -649,35 +649,6 @@ def __call__(self, stream, content_type):
             stream.close()
 
 
-class StringDeserializer(object):
-    """Return the response as a decoded string.
-
-    Args:
-        encoding (str): The string encoding to use (default=utf-8).
-        accept (str): The Accept header to send to the server (optional).
-    """
-
-    def __init__(self, encoding="utf-8", accept=None):
-        """
-        Args:
-            encoding:
-            accept:
-        """
-        self.encoding = encoding
-        self.accept = accept
-
-    def __call__(self, stream, content_type):
-        """
-        Args:
-            stream:
-            content_type:
-        """
-        try:
-            return stream.read().decode(self.encoding)
-        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

src/sagemaker/test_deserializers.py

@@ -0,0 +1,33 @@
+# Copyright 2017-2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+from __future__ import absolute_import
+
+import io
+
+from sagemaker.deserializers import StringDeserializer
+
+
+def test_string_deserializer_plain_text():
+    deserializer = StringDeserializer()
+
+    result = deserializer.deserialize("Hello, world!", "text/plain")
+
+    assert result == "Hello, world!"
+
+
+def test_string_deserializer_octet_stream():
+    deserializer = StringDeserializer()
+
+    result = deserializer.deserialize(io.BytesIO(b"Hello, world!"), "application/octet-stream")
+
+    assert result == "Hello, world!"

src/sagemaker/utils.py

@@ -745,3 +745,23 @@ def _module_import_error(py_module, feature, extras):
         "to install all required dependencies."
     )
     return error_msg.format(py_module, feature, extras)
+
+
+def parse_mime_type(mime_type):
+    """Parse a MIME type and return the type, subtype, and parameters.
+
+    Args:
+        mime_type (str): A MIME type.
+
+    Returns:
+        tuple: A three-tuple containing the type, subtype, and parameters. The
+            type and subtype are strings, and the parameters are stored in a
+            dictionary.
+    """
+    category, remaining = mime_type.split("/")
+    subtype = remaining.split(";")[0]
+    parameters = {}
+    for parameter in remaining.split(";")[1:]:
+        attribute, value = parameter.split("=")
+        parameters[attribute] = value
+    return category, subtype, parameters

tests/integ/test_multidatamodel.py

@@ -24,9 +24,10 @@
 
 from sagemaker import utils
 from sagemaker.amazon.randomcutforest import RandomCutForest
+from sagemaker.deserializers import StringDeserializer
 from sagemaker.multidatamodel import MultiDataModel
 from sagemaker.mxnet import MXNet
-from sagemaker.predictor import Predictor, StringDeserializer, npy_serializer
+from sagemaker.predictor import Predictor, npy_serializer
 from sagemaker.utils import sagemaker_timestamp, unique_name_from_base, get_ecr_image_uri_prefix
 from tests.integ import DATA_DIR, TRAINING_DEFAULT_TIMEOUT_MINUTES
 from tests.integ.retry import retries

tests/unit/test_predictor.py

@@ -27,7 +27,6 @@
     csv_serializer,
     csv_deserializer,
     BytesDeserializer,
-    StringDeserializer,
     StreamDeserializer,
     numpy_deserializer,
     npy_serializer,
@@ -190,12 +189,6 @@ def test_bytes_deserializer():
     assert result == b"[1, 2, 3]"
 
 
-def test_string_deserializer():
-    result = StringDeserializer()(io.BytesIO(b"[1, 2, 3]"), "application/json")
-
-    assert result == "[1, 2, 3]"
-
-
 def test_stream_deserializer():
     stream, content_type = StreamDeserializer()(io.BytesIO(b"[1, 2, 3]"), "application/json")
     result = stream.read()

tests/unit/test_utils.py

@@ -779,3 +779,11 @@ def test_partition_by_region():
     assert sagemaker.utils._aws_partition("us-gov-east-1") == "aws-us-gov"
     assert sagemaker.utils._aws_partition("us-iso-east-1") == "aws-iso"
     assert sagemaker.utils._aws_partition("us-isob-east-1") == "aws-iso-b"
+
+
+def test_parse_mime_type():
+    mime_type = "application/octet-stream;charset=UTF-8"
+    category, subtype, parameters = sagemaker.utils.parse_mime_type(mime_type)
+    assert category == "application"
+    assert subtype == "octet-stream"
+    assert parameters == {"charset": "UTF-8"}