#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Fri Jan 8 15:49:55 2021
@author: ikaheia1
Function wrapper for handling the plotting of figure or saving it to the
destined folder.
"""
import matplotlib.pyplot as plt
import functools
[docs]def plot_decorator(func):
""" Save the plot figure if savename and path are in function
keyword arguments.
Parameters
----------
func : function
Plotting function to be decorated.
Raises
------
Exception
Requested folder does not exist og the agruments are not given in
correct way.
Returns
-------
wrapper : function wrapper
Wrapped plotting function.
"""
@functools.wraps(func)
def wrapper(*args,**kwargs):
""" Function wrapper for handling the plots.
Parameters
----------
*args : arguments
Arguments of the wrapped function
**kwargs : keyword arguments
Keyword arguments of the wrapped function
Raises
------
Exception
Requested folder does not exist og the agruments are not given in
correct way.
Returns
-------
fig : matplotlib plt object
Plotted figure
"""
fig = func(*args,**kwargs)
sn = ()
sp = ()
if 'savename' in kwargs:
sn = kwargs.get('savename')
if 'savepath' in kwargs:
sp = kwargs.get('savepath')
if not all((sn,sp)):
plt.show()
elif all((sn,sp)):
assert isinstance(sn,str), "Invalid savename type."
if sp.exists():
with open(sp / (sn + ".png"), mode="wb") as outfile:
#plt.close()
plt.savefig(outfile, format="png")
else:
raise Exception("Requested folder: " + str(sp) + " does not exist.")
else:
raise Exception("Arguments were not given correctly.")
if 'test' in kwargs:
if kwargs.get('test') == True:
return fig
return wrapper