107 lines
3.9 KiB
Python
107 lines
3.9 KiB
Python
|
import io
|
||
|
|
import ntpath
|
||
|
|
import base64
|
||
|
|
|
||
|
|
# region Utils for Download component
|
||
|
|
|
||
|
|
|
||
|
|
def send_file(path, filename=None, type=None):
|
||
|
|
"""
|
||
|
|
Convert a file into the format expected by the Download component.
|
||
|
|
:param path: path to the file to be sent
|
||
|
|
:param filename: name of the file, if not provided the original filename is used
|
||
|
|
:param type: type of the file (optional, passed to Blob in the javascript layer)
|
||
|
|
:return: dict of file content (base64 encoded) and meta data used by the Download component
|
||
|
|
"""
|
||
|
|
# If filename is not set, read it from the path.
|
||
|
|
if filename is None:
|
||
|
|
filename = ntpath.basename(path)
|
||
|
|
# Read the file contents and send it.
|
||
|
|
with open(path, "rb") as f:
|
||
|
|
return send_bytes(f.read(), filename, type)
|
||
|
|
|
||
|
|
|
||
|
|
def send_bytes(src, filename, type=None, **kwargs):
|
||
|
|
"""
|
||
|
|
Convert data written to BytesIO into the format expected by the Download component.
|
||
|
|
:param src: array of bytes or a writer that can write to BytesIO
|
||
|
|
:param filename: the name of the file
|
||
|
|
:param type: type of the file (optional, passed to Blob in the javascript layer)
|
||
|
|
:return: dict of data frame content (base64 encoded) and meta data used by the Download component
|
||
|
|
"""
|
||
|
|
content = src if isinstance(src, bytes) else _io_to_str(io.BytesIO(), src, **kwargs)
|
||
|
|
return dict(
|
||
|
|
content=base64.b64encode(content).decode(),
|
||
|
|
filename=filename,
|
||
|
|
type=type,
|
||
|
|
base64=True,
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def send_string(src, filename, type=None, **kwargs):
|
||
|
|
"""
|
||
|
|
Convert data written to StringIO into the format expected by the Download component.
|
||
|
|
:param src: a string or a writer that can write to StringIO
|
||
|
|
:param filename: the name of the file
|
||
|
|
:param type: type of the file (optional, passed to Blob in the javascript layer)
|
||
|
|
:return: dict of data frame content (NOT base64 encoded) and meta data used by the Download component
|
||
|
|
"""
|
||
|
|
content = src if isinstance(src, str) else _io_to_str(io.StringIO(), src, **kwargs)
|
||
|
|
return dict(content=content, filename=filename, type=type, base64=False)
|
||
|
|
|
||
|
|
|
||
|
|
def _io_to_str(data_io, writer, **kwargs):
|
||
|
|
# Some pandas writers try to close the IO, we do not want that.
|
||
|
|
data_io_close = data_io.close
|
||
|
|
data_io.close = lambda: None
|
||
|
|
# Write data content.
|
||
|
|
writer(data_io, **kwargs)
|
||
|
|
data_value = data_io.getvalue()
|
||
|
|
data_io_close()
|
||
|
|
return data_value
|
||
|
|
|
||
|
|
|
||
|
|
def send_data_frame(writer, filename, type=None, **kwargs):
|
||
|
|
"""
|
||
|
|
Convert data frame into the format expected by the Download component.
|
||
|
|
:param writer: a data frame writer
|
||
|
|
:param filename: the name of the file
|
||
|
|
:param type: type of the file (optional, passed to Blob in the javascript layer)
|
||
|
|
:return: dict of data frame content (base64 encoded) and meta data used by the Download component
|
||
|
|
|
||
|
|
Examples
|
||
|
|
--------
|
||
|
|
|
||
|
|
>>> df = pd.DataFrame({'a': [1, 2, 3, 4], 'b': [2, 1, 5, 6], 'c': ['x', 'x', 'y', 'y']})
|
||
|
|
...
|
||
|
|
>>> send_data_frame(df.to_csv, "mydf.csv") # download as csv
|
||
|
|
>>> send_data_frame(df.to_json, "mydf.json") # download as json
|
||
|
|
>>> send_data_frame(df.to_excel, "mydf.xls", index=False) # download as excel
|
||
|
|
>>> send_data_frame(df.to_pickle, "mydf.pkl") # download as pickle
|
||
|
|
|
||
|
|
"""
|
||
|
|
name = writer.__name__
|
||
|
|
# Check if the provided writer is known.
|
||
|
|
if name not in _data_frame_senders.keys():
|
||
|
|
raise ValueError(
|
||
|
|
"The provided writer ({}) is not supported, "
|
||
|
|
"try calling send_string or send_bytes directly.".format(name)
|
||
|
|
)
|
||
|
|
# Send data frame using the appropriate send function.
|
||
|
|
return _data_frame_senders[name](writer, filename, type, **kwargs)
|
||
|
|
|
||
|
|
|
||
|
|
_data_frame_senders = {
|
||
|
|
"to_csv": send_string,
|
||
|
|
"to_json": send_string,
|
||
|
|
"to_html": send_string,
|
||
|
|
"to_excel": send_bytes,
|
||
|
|
"to_feather": send_bytes,
|
||
|
|
"to_parquet": send_bytes,
|
||
|
|
"to_msgpack": send_bytes,
|
||
|
|
"to_stata": send_bytes,
|
||
|
|
"to_pickle": send_bytes,
|
||
|
|
}
|
||
|
|
|
||
|
|
# endregion
|