diff --git a/src/py/formatters/pickle.py b/src/py/formatters/pickle.py index 62693e24..46fcf6a4 100644 --- a/src/py/formatters/pickle.py +++ b/src/py/formatters/pickle.py @@ -16,6 +16,11 @@ class PickleFormatter(BaseFormatter): decode_format = "json" def decode(self, value): + def get_json_output(deserialized_object): + return json.dumps(deserialized_object, + default=self.default, + ensure_ascii=False) + read_only = self.read_only decode_format = self.decode_format deserialized = '' @@ -43,10 +48,11 @@ class PickleFormatter(BaseFormatter): border=0) output = self.format_html_output(deserialized, html) decode_format = 'html' + else: + output = get_json_output(deserialized) else: - output = json.dumps(deserialized, - default=self.default, - ensure_ascii=False) + output = get_json_output(deserialized) + return { 'output': output, 'decode-format': decode_format, diff --git a/tests/py_tests/test_formatters/test_pickle_formatter.py b/tests/py_tests/test_formatters/test_pickle_formatter.py index 4de754e6..136fbcb7 100644 --- a/tests/py_tests/test_formatters/test_pickle_formatter.py +++ b/tests/py_tests/test_formatters/test_pickle_formatter.py @@ -30,10 +30,19 @@ class TestPickleFormatter(unittest.TestCase): columns=list('ABCD')), ) def test_decode(self, val): - pickled_val = pickle.dumps(val) - expected_output = json.dumps(val, default=self.formatter.default, - ensure_ascii=False) + if isinstance(val, pd.Series): + expected_output = f'{str(type(val))[1:-1]}\n{val.to_string()}' + elif isinstance(val, pd.DataFrame): + html = val.to_html(render_links=True, border=0) + expected_output = self.formatter.format_html_output(val, html) + elif isinstance(val, np.ndarray): + html = pd.DataFrame(val).to_html(render_links=True, border=0) + expected_output = self.formatter.format_html_output(val, html) + else: + expected_output = json.dumps(val, default=self.formatter.default, + ensure_ascii=False) + pickled_val = pickle.dumps(val) formatter_response_dict = self.formatter.decode(pickled_val) self.assertIn('output', formatter_response_dict)