"""
Main uses: save and cache data, save figures, change figure color, timing.
"""
import inspect
from pathlib import Path
from functools import wraps
from collections.abc import Iterable
from importlib.util import find_spec
from ._helper import NoFigure
if find_spec("plotly") is None:
plotly_figure = NoFigure
else:
from plotly.graph_objs._figure import Figure as plotly_figure
if find_spec("matplotlib") is None:
mpl_figure = NoFigure
else:
from matplotlib.figure import Figure as mpl_figure
import matplotlib.pyplot as plt
from . import storage, config
from .paths import datapath, figpath
from .inspection import classify_call_attrs, merge_wrapper_signatures
from ._helper import merge_nested_dict
[docs]def savedata(keys_or_function=None, include_classes="file",
ext=config.EXT_DEFAULT_DATA, keys=config.KEYS_DEFAULT_DATA, funcname_in_filename=config.FUNCNAME_IN_FILENAME_DEFAULT_DATA,
overwrite=False, save=True, load_opts_default_save=True, #defaults for extra arguments
load_opts={}, **save_opts):
"""
Decorator for automatically saving output and then loading cached data.
Default behavior:
1st function call: stores the data with extension 'ext'.
rest: loads stored data if the key args passed to the function coincide.
NOTE (!) decorated funcs will have extra arguments:
- save (bool). whether to save the output of the function.
- overwrite (bool). whether to overwrite the saved version.
- funcname_in_filename(bool) whether to include funcname in the filename.
- keys (dict/str/iterable) · dict: specifies filename of the form k1-v1_k2-v2_...kn_vn.
k_i do not have to be arguments of the function.
· str: key of a keyword argument. Example: keys = 'x'.
"kwargs_defaults" => default kwargs that were not modified.
"kwargs" => kwargs passed during function call.
"kwargs_full" => kws_defaults + kws.
"pos_only" => length of *args
"args" => attrs != **kwargs, *args.
"all" if config.KEYS_ADD_POSONLY_TO_ALL => all attrs
else => all attrs except pos_only: kwargs_full + args
can combine options using "+" or "-". Example: "args+z", "x+y+kwargs", "all-y".
· iterable: containing a combination of the available string keys (above). ["k1", "k2"] == "k1+k2".
Attrs:
- function: function to which the decorator is applied
- ext: storing extension. Selects 'storage' functions save_ext, load_ext.
Supported: 'lzma' (default), 'bz2', 'json', 'csv', 'npz'.
- include_classes: include class tree in saving_path.
- load_opts: kws for storage.load_ext. default kws are those of 'saving_options', those specified update saving_options dict.
- save_opts: kws for storage.save_ext
- load_opts_default_save: use save_opts as default for load_opts.
- rest: default behavior for decorated funcs extra arguments (above).
Returns: Function decorator
"""
if isinstance(keys_or_function, Iterable):
keys = keys_or_function
func = None
elif callable(keys_or_function):
func = keys_or_function
else:
func = None
if load_opts_default_save:
load_opts = {**save_opts, **load_opts}
def _savedata(func):
@wraps(func)
def wrapper(*args, overwrite=overwrite, keys=keys, save=save, funcname_in_filename=funcname_in_filename, **kwargs):
key_opts = classify_call_attrs(func, args, kwargs, add_pos_only_to_all=config.KEYS_ADD_POSONLY_TO_ALL)
save_keys = merge_nested_dict(key_opts, keys, key_default="all")
saving_path = datapath(keys=save_keys, func=func, ext=ext, include_classes=include_classes, funcname_in_filename=funcname_in_filename)
if Path(saving_path).exists() and not overwrite:
return getattr(storage, f"load_{ext}")(saving_path, **load_opts)
else:
result = func(*args, **kwargs)
getattr(storage, f"save_{ext}")(result, saving_path, **save_opts)
return result
wrapper.__signature__ = merge_wrapper_signatures(wrapper, ["overwrite", "keys", "save", "funcname_in_filename"])
wrapper.__out__ = "data"
return wrapper
if func is None:
return _savedata
else:
return _savedata(func)
[docs]def savefig(keys_or_function=None, include_classes="file",
ext=config.EXT_DEFAULT_FIG, keys=config.KEYS_DEFAULT_FIG, funcname_in_filename=config.FUNCNAME_IN_FILENAME_DEFAULT_FIG, return_fig=config.RETURN_FIG_DEFAULT,
overwrite=True, save=True, #defaults for extra arguments
**save_opts):
"""
Generates figsaver decorator.
Figure returned by function is automatically saved. Compatible with matplotlib and plotly.
NOTE (!) decorated funcs will have extra arguments:
- return_fig (bool) whether to return the figure (output of function).
- save (bool). whether to save the figure.
- overwrite (bool). whether to overwrite the saved version.
- funcname_in_filename(bool) whether to include funcname in the filename.
- keys (dict/str/iterable) · dict: specifies filename of the form k1-v1_k2-v2_...kn_vn.
k_i do not have to be arguments of the function.
· str: key of a keyword argument. Example: keys = 'x'.
"kwargs_defaults" => default kwargs that were not modified.
"kwargs" => kwargs passed during function call.
"kwargs_full" => kws_defaults + kws.
"pos_only" => length of *args. Also self and cls are counted as pos_only arguments.
"args" => attrs != **kwargs, *args.
"all" if config.KEYS_ADD_POSONLY_TO_ALL => all attrs
else => all attrs except pos_only: kwargs_full + args
can combine options using "+" or "-". Example: "args+z", "x+y+kwargs", "all-y".
· iterable: containing a combination of the available string keys (above). ["k1", "k2"] == "k1+k2".
Attrs:
- function: function to which the decorator is applied
- ext: storing extension. 'eps' recommended for articles.
Supported: any extension supported by matplotlib/plotly. Example: 'png', 'eps', 'html' (plotly), etc.
- include_classes: include class tree in saving_path.
- save_opts: kws for saving function.
- rest: default behavior for decorated funcs extra arguments (above).
Returns: Function decorator
"""
if isinstance(keys_or_function, Iterable):
keys = keys_or_function
func = None
elif callable(keys_or_function):
func = keys_or_function
else:
func = None
mpl_save_defaults = dict(bbox_inches="tight")
def _savefig(func):
@wraps(func)
def wrapper(*args, overwrite=overwrite, keys=keys, save=save, return_fig=return_fig, funcname_in_filename=funcname_in_filename, **kwargs):
fig = func(*args, **kwargs)
key_opts = classify_call_attrs(func, args, kwargs, add_pos_only_to_all=config.KEYS_ADD_POSONLY_TO_ALL)
save_keys = merge_nested_dict(key_opts, keys, key_default="all")
saving_path = figpath(keys=save_keys, func=func, ext=ext, include_classes=include_classes, funcname_in_filename=funcname_in_filename)
if not Path(saving_path).exists() or overwrite:
if isinstance(fig, mpl_figure):
fig.savefig(saving_path, format=ext, **{**mpl_save_defaults, **save_opts})
plt.close(fig)
elif isinstance(fig, plotly_figure):
if ext == "html":
fig.write_html(saving_path, **save_opts)
else:
fig.write_image(saving_path, format=ext, **save_opts)
else:
raise TypeError(f"fig type '{type(fig)}' not valid. Available: 'matplotlib.figure.Figure', 'plotly.grap_objs._figure.Figure'.")
if return_fig:
return fig
else:
return
wrapper.__signature__ = merge_wrapper_signatures(wrapper, ["overwrite", "keys", "save", "funcname_in_filename", "return_fig"])
wrapper.__out__ = "figure"
return wrapper
if func is None:
return _savefig
else:
return _savefig(func)