def write_output(output_format, output, output_file):
"""
Write output to output_file using the format from output_format.
If output_format == None, write output as pickle without indicating format (legacy method)
:param output_format:
:param output:
:param output_file:
:return:
"""
with open(output_file, 'wb') as fp:
if output_format:
# Indicate output format
fp.write(output_format.encode() + b'.')
# Write actual data
output_format = DataFormat(output_format.lower())
serialized = serialization.serialize(output, output_format)
fp.write(serialized)
else:
# No output format specified, use legacy method
fp.write(pickle.dumps(output))
def test_pickle_serialization_pandas():
data = pd.DataFrame([1, 2, 3])
pickled = serialization.serialize(data, DataFormat.PICKLE)
pd.testing.assert_frame_equal(data, pickle.loads(pickled))
def test_pickle_serialization(data):
pickled = serialization.serialize(data, DataFormat.PICKLE)
assert data == pickle.loads(pickled)
def test_json_serialization(data, target):
result = serialization.serialize(data, DataFormat.JSON)
assert target == result.decode()