def as_bytes(self, data, extension=None):
if extension is None:
extension = self.default_extension()
assert self.is_type_of(data)
mimetype = mimetype_from_extension(extension)
if extension == "csv":
output = StringIO()
data.to_csv(output, index=False)
return output.getvalue().encode("utf-8"), mimetype
elif extension == "tsv":
output = StringIO()
data.to_csv(output, index=False, sep="\t")
return output.getvalue().encode("utf-8"), mimetype
elif extension == "json":
output = StringIO()
data.to_json(output, index=False, orient="table")
return output.getvalue().encode("utf-8"), mimetype
elif extension in ("html", "htm"):
output = StringIO()
data.to_html(output, index=False)
return output.getvalue().encode("utf-8"), mimetype
elif extension in ("pkl", "pickle"):
output = ResilientBytesIO()
data.to_pickle(output, compression=None)
b = output.getvalue()
output.really_close()
return b, mimetype
elif extension == "parquet":
output = ResilientBytesIO()
data.to_parquet(output, engine="pyarrow")
b = output.getvalue()
output.really_close()
return b, mimetype
elif extension == "feather":
output = ResilientBytesIO()
data.to_feather(output)
b = output.getvalue()
output.really_close()
return b, mimetype
elif extension == "xlsx":
output = BytesIO()
writer = pd.ExcelWriter(output, engine="xlsxwriter")
data.to_excel(writer)
writer.close()
return output.getvalue(), mimetype
elif extension == "msgpack":
output = BytesIO()
data.to_msgpack(output)
return output.getvalue(), mimetype
else:
raise Exception(
f"Serialization: file extension {extension} is not supported by dataframe type."
)
