Merge branch 'zwei' into add-pandas-deserializer

Author: bveeramani

src/sagemaker/deserializers.py

@@ -38,11 +38,11 @@ class BaseDeserializer(abc.ABC):
     """
 
     @abc.abstractmethod
-    def deserialize(self, data, content_type):
+    def deserialize(self, stream, content_type):
         """Deserialize data received from an inference endpoint.
 
         Args:
-            data (object): Data to be deserialized.
+            stream (botocore.response.StreamingBody): Data to be deserialized.
             content_type (str): The MIME type of the data.
 
         Returns:
@@ -68,41 +68,41 @@ def __init__(self, encoding="UTF-8"):
         """
         self.encoding = encoding
 
-    def deserialize(self, data, content_type):
+    def deserialize(self, stream, content_type):
         """Deserialize data from an inference endpoint into a decoded string.
 
         Args:
-            data (object): Data to be deserialized.
+            stream (botocore.response.StreamingBody): Data to be deserialized.
             content_type (str): The MIME type of the data.
 
         Returns:
             str: The data deserialized into a decoded string.
         """
         try:
-            return data.read().decode(self.encoding)
+            return stream.read().decode(self.encoding)
         finally:
-            data.close()
+            stream.close()
 
 
 class BytesDeserializer(BaseDeserializer):
     """Deserialize a stream of bytes into a bytes object."""
 
     ACCEPT = "*/*"
 
-    def deserialize(self, data, content_type):
+    def deserialize(self, stream, content_type):
         """Read a stream of bytes returned from an inference endpoint.
 
         Args:
-            data (object): A stream of bytes.
+            stream (botocore.response.StreamingBody): 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()
+            return stream.read()
         finally:
-            data.close()
+            stream.close()
 
 
 class CSVDeserializer(BaseDeserializer):
@@ -118,22 +118,22 @@ def __init__(self, encoding="utf-8"):
         """
         self.encoding = encoding
 
-    def deserialize(self, data, content_type):
+    def deserialize(self, stream, content_type):
         """Deserialize data from an inference endpoint into a list of lists.
 
         Args:
-            data (botocore.response.StreamingBody): Data to be deserialized.
+            stream (botocore.response.StreamingBody): Data to be deserialized.
             content_type (str): The MIME type of the data.
 
         Returns:
             list: The data deserialized into a list of lists representing the
                 contents of a CSV file.
         """
         try:
-            decoded_string = data.read().decode(self.encoding)
+            decoded_string = stream.read().decode(self.encoding)
             return list(csv.reader(decoded_string.splitlines()))
         finally:
-            data.close()
+            stream.close()
 
 
 class StreamDeserializer(BaseDeserializer):
@@ -145,17 +145,17 @@ class StreamDeserializer(BaseDeserializer):
 
     ACCEPT = "*/*"
 
-    def deserialize(self, data, content_type):
+    def deserialize(self, stream, content_type):
         """Returns a stream of the response body and the MIME type of the data.
 
         Args:
-            data (object): A stream of bytes.
+            stream (botocore.response.StreamingBody): A stream of bytes.
             content_type (str): The MIME type of the data.
 
         Returns:
             tuple: A two-tuple containing the stream and content-type.
         """
-        return data, content_type
+        return stream, content_type
 
 
 class NumpyDeserializer(BaseDeserializer):
@@ -171,11 +171,11 @@ def __init__(self, dtype=None):
         """
         self.dtype = dtype
 
-    def deserialize(self, data, content_type):
+    def deserialize(self, stream, content_type):
         """Deserialize data from an inference endpoint into a NumPy array.
 
         Args:
-            data (botocore.response.StreamingBody): Data to be deserialized.
+            stream (botocore.response.StreamingBody): Data to be deserialized.
             content_type (str): The MIME type of the data.
 
         Returns:
@@ -184,14 +184,14 @@ def deserialize(self, data, content_type):
         try:
             if content_type == "text/csv":
                 return np.genfromtxt(
-                    codecs.getreader("utf-8")(data), delimiter=",", dtype=self.dtype
+                    codecs.getreader("utf-8")(stream), delimiter=",", dtype=self.dtype
                 )
             if content_type == "application/json":
-                return np.array(json.load(codecs.getreader("utf-8")(data)), dtype=self.dtype)
+                return np.array(json.load(codecs.getreader("utf-8")(stream)), dtype=self.dtype)
             if content_type == "application/x-npy":
-                return np.load(io.BytesIO(data.read()))
+                return np.load(io.BytesIO(stream.read()))
         finally:
-            data.close()
+            stream.close()
 
         raise ValueError("%s cannot read content type %s." % (__class__.__name__, content_type))
 
@@ -201,45 +201,45 @@ class JSONDeserializer(BaseDeserializer):
 
     ACCEPT = "application/json"
 
-    def deserialize(self, data, content_type):
+    def deserialize(self, stream, content_type):
         """Deserialize JSON data from an inference endpoint into a Python object.
 
         Args:
-            data (botocore.response.StreamingBody): Data to be deserialized.
+            stream (botocore.response.StreamingBody): Data to be deserialized.
             content_type (str): The MIME type of the data.
 
         Returns:
             object: The JSON-formatted data deserialized into a Python object.
         """
         try:
-            return json.load(codecs.getreader("utf-8")(data))
+            return json.load(codecs.getreader("utf-8")(stream))
         finally:
-            data.close()
+            stream.close()
 
 
 class PandasDeserializer(BaseDeserializer):
     """Deserialize CSV or JSON data from an inference endpoint into a pandas dataframe."""
 
     ACCEPT = "text/csv"
 
-    def deserialize(self, data, content_type):
+    def deserialize(self, stream, content_type):
         """Deserialize CSV or JSON data from an inference endpoint into a pandas
         dataframe.
 
         If the data is JSON, the data should be formatted in the 'columns' orient.
         See https://pandas.pydata.org/pandas-docs/stable/reference/api/pandas.read_json.html
 
         Args:
-            data (botocore.response.StreamingBody): Data to be deserialized.
+            stream (botocore.response.StreamingBody): Data to be deserialized.
             content_type (str): The MIME type of the data.
 
         Returns:
             pandas.DataFrame: The data deserialized into a pandas DataFrame.
         """
         if content_type == "text/csv":
-            return pandas.read_csv(data)
+            return pandas.read_csv(stream)
 
         if content_type == "application/json":
-            return pandas.read_json(data)
+            return pandas.read_json(stream)
 
         raise ValueError("%s cannot read content type %s." % (__class__.__name__, content_type))