Refactor code

Author: Balaji Veeramani

src/sagemaker/deserializers.py

@@ -15,8 +15,6 @@
 
 import abc
 
-from sagemaker.utils import parse_mime_type
-
 
 class BaseDeserializer(abc.ABC):
     """Abstract base class for creation of new deserializers.
@@ -43,35 +41,30 @@ def ACCEPT(self):
         """The content type that is expected from the inference endpoint."""
 
 
-class StringDeserializer(object):
+class StringDeserializer(BaseDeserializer):
     """Deserialize data from an inference endpoint into a decoded string."""
 
+    ACCEPT = "application/json"
+
     def __init__(self, encoding="UTF-8"):
-        """Initialize the default encoding.
+        """Initialize the string encoding.
 
         Args:
-            encoding (str): The string encoding to use, if a charset is not
-                provided by the server (default: UTF-8).
+            encoding (str): The string encoding to use (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.
+            data (object): Data to be deserialized.
             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)
+            return data.read().decode(self.encoding)
         finally:
             data.close()

src/sagemaker/utils.py

@@ -745,23 +745,3 @@ 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/unit/sagemaker/test_deserializers.py

@@ -17,17 +17,9 @@
 from sagemaker.deserializers import StringDeserializer
 
 
-def test_string_deserializer_plain_text():
+def test_string_deserializer():
     deserializer = StringDeserializer()
 
-    result = deserializer.deserialize("Hello, world!", "text/plain")
+    result = deserializer.deserialize(io.BytesIO(b"[1, 2, 3]"), "application/json")
 
-    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!"
+    assert result == "[1, 2, 3]"

tests/unit/test_utils.py

@@ -779,11 +779,3 @@ 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"}