Fixed can_encode for JSONCodec and a pandas test

Author: RobertSamoilescu

mlserver/codecs/json.py

@@ -78,6 +78,23 @@ def encode_to_json(v: Any, use_bytes: bool = True) -> Union[str, bytes]:
     return enc_v
 
 
+def _is_primitive(obj):
+    return isinstance(obj, (int, float, str, bool, type(None)))
+
+
+def _is_nested_primitives(obj):
+    if _is_primitive(obj):
+        return True
+    elif isinstance(obj, list):
+        return all(_is_nested_primitives(item) for item in obj)
+    elif isinstance(obj, dict):
+        return all(
+            isinstance(key, str) and _is_nested_primitives(value)
+            for key, value in obj.items()
+        )
+    return False
+
+
 @register_input_codec
 class JSONCodec(InputCodec):
     """
@@ -89,12 +106,23 @@ class JSONCodec(InputCodec):
 
     @classmethod
     def can_encode(cls, payload: Any) -> bool:
-        try:
-            encode_to_json(payload)
-            return True
-        except Exception:
+        is_json = all(_is_nested_primitives(item) for item in payload)
+
+        if not is_json:
             return False
 
+        all_primitive = all(_is_primitive(item) for item in payload)
+
+        # have to do it this way in case payload is not indexable
+        types = [type(item) for item in payload]
+        same_type = [types[0] == item for item in types]
+
+        # Don't want to json encode a list of primitives of the same type
+        if all_primitive and same_type:
+            return False
+
+        return True
+
     @classmethod
     def encode_output(
         cls, name: str, payload: List[Any], use_bytes: bool = True, **kwargs

tests/codecs/test_json.py

@@ -49,10 +49,17 @@ def test_encode_input(input: Any, expected: bytes):
 @pytest.mark.parametrize(
     "payload, expected",
     [
-        ([1, 2, 3], True),
-        ({"dummy_1": 1, "dummy_2": 2}, True),
-        (np.array([1, 2, 3]), False),
-        (set([1, 2, 3]), False),
+        ([[1, 2, 3]], True),
+        ([{"dummy_1": 1, "dummy_2": 2}], True),
+        (
+            [
+                {"dummy_1": 1},
+                {"dummy_2": [{"dummy_3": 3, "dummy_4": 4}, {"dummy_5": 5}]},
+            ],
+            True,
+        ),
+        ([np.array([1, 2, 3])], False),
+        ([{1, 2, 3}], False),
     ],
 )
 def test_json_codec_can_encode(payload: Any, expected: bool):

tests/codecs/test_pandas.py

@@ -6,6 +6,7 @@
 
 from mlserver.codecs.pandas import PandasCodec, _to_response_output
 from mlserver.codecs.string import StringCodec
+from mlserver.codecs.json import JSONCodec
 from mlserver.types import (
     InferenceRequest,
     InferenceResponse,
@@ -67,7 +68,11 @@ def test_can_encode(payload: Any, expected: bool):
             pd.Series(data=[[1, 2, 3], [4, 5, 6]], name="bar"),
             True,
             ResponseOutput(
-                name="bar", shape=[2, 1], data=[[1, 2, 3], [4, 5, 6]], datatype="BYTES"
+                name="bar",
+                shape=[2, 1],
+                data=[b"[1,2,3]", b"[4,5,6]"],
+                datatype="BYTES",
+                parameters=Parameters(content_type=JSONCodec.ContentType),
             ),
         ),
         (