Merge branch 'zwei' into model-dir-slashes

Author: laurenyu

src/sagemaker/serializers.py

@@ -14,6 +14,7 @@
 from __future__ import absolute_import
 
 import abc
+from collections.abc import Iterable
 import csv
 import io
 import json
@@ -192,6 +193,34 @@ def serialize(self, data):
         return json.dumps(data)
 
 
+class JSONLinesSerializer(BaseSerializer):
+    """Serialize data to a JSON Lines formatted string."""
+
+    CONTENT_TYPE = "application/jsonlines"
+
+    def serialize(self, data):
+        """Serialize data of various formats to a JSON Lines formatted string.
+
+        Args:
+            data (object): Data to be serialized. The data can be a string,
+                iterable of JSON serializable objects, or a file-like object.
+
+        Returns:
+            str: The data serialized as a string containing newline-separated
+                JSON values.
+        """
+        if isinstance(data, str):
+            return data
+
+        if hasattr(data, "read"):
+            return data.read()
+
+        if isinstance(data, Iterable):
+            return "\n".join(json.dumps(element) for element in data)
+
+        raise ValueError("Object of type %s is not JSON Lines serializable." % type(data))
+
+
 class SparseMatrixSerializer(BaseSerializer):
     """Serialize a sparse matrix to a buffer using the .npz format."""
 

tests/unit/sagemaker/test_serializers.py

@@ -25,6 +25,7 @@
     NumpySerializer,
     JSONSerializer,
     SparseMatrixSerializer,
+    JSONLinesSerializer,
 )
 from tests.unit import DATA_DIR
 
@@ -235,6 +236,57 @@ def test_json_serializer_csv_buffer(json_serializer):
         assert result == validation_value
 
 
+@pytest.fixture
+def json_lines_serializer():
+    return JSONLinesSerializer()
+
+
+@pytest.mark.parametrize(
+    "input, expected",
+    [
+        ('["Name", "Score"]\n["Gilbert", 24]', '["Name", "Score"]\n["Gilbert", 24]'),
+        (
+            '{"Name": "Gilbert", "Score": 24}\n{"Name": "Alexa", "Score": 29}',
+            '{"Name": "Gilbert", "Score": 24}\n{"Name": "Alexa", "Score": 29}',
+        ),
+    ],
+)
+def test_json_lines_serializer_string(json_lines_serializer, input, expected):
+    actual = json_lines_serializer.serialize(input)
+    assert actual == expected
+
+
+@pytest.mark.parametrize(
+    "input, expected",
+    [
+        ([["Name", "Score"], ["Gilbert", 24]], '["Name", "Score"]\n["Gilbert", 24]'),
+        (
+            [{"Name": "Gilbert", "Score": 24}, {"Name": "Alexa", "Score": 29}],
+            '{"Name": "Gilbert", "Score": 24}\n{"Name": "Alexa", "Score": 29}',
+        ),
+    ],
+)
+def test_json_lines_serializer_list(json_lines_serializer, input, expected):
+    actual = json_lines_serializer.serialize(input)
+    assert actual == expected
+
+
+@pytest.mark.parametrize(
+    "source, expected",
+    [
+        ('["Name", "Score"]\n["Gilbert", 24]', '["Name", "Score"]\n["Gilbert", 24]'),
+        (
+            '{"Name": "Gilbert", "Score": 24}\n{"Name": "Alexa", "Score": 29}',
+            '{"Name": "Gilbert", "Score": 24}\n{"Name": "Alexa", "Score": 29}',
+        ),
+    ],
+)
+def test_json_lines_serializer_file_like(json_lines_serializer, source, expected):
+    input = io.StringIO(source)
+    actual = json_lines_serializer.serialize(input)
+    assert actual == expected
+
+
 @pytest.fixture
 def sparse_matrix_serializer():
     return SparseMatrixSerializer()