use python decorator to automatically replace function argument default value?

Question

Actually the title does not reflect exactly the question I want to ask. My purpose is like this: I am writing some plotting functions by using matplotlib. I have a series of functions that are for different plotting purpose. like line_plot() for lines, bar_plot() for bar etc. for example:

import matplotlib.pyplot as plt
def line_plot(axes=None,x=None,y=None):
    if axes==None:
        fig=plt.figure()
        axes=fig.add_subplot(111)
    else:
        pass
    axes.plot(x,y)

def bar_plot(axes=None,x=None,y=None):
    if axes==None:
        fig=plt.figure()
        axes=fig.add_subplot(111)
    else:
        pass
    axes.bar(left=x,height=y)

Yet the problem is, for each function that's defined, I have to repeat this part of code:

    if axes==None:
        fig=plt.figure()
        axes=fig.add_subplot(111)
    else:
        pass

Is there way like by using a decorator, which I can apply before the definition of the plotting function, that will do the repeated part of code automatically? Thus I don't have to repeat them each time.

one possible option is to define a function like this:

def check_axes(axes):
    if axes==None:
        fig=plt.figure()
        axes=fig.add_subplot(111)
        return axes
    else:
        return axes

Then the examples will be like:

import matplotlib.pyplot as plt    
def line_plot(axes=None,x=None,y=None):
    axes=check_axes(axes)
    axes.plot(x,y)

def bar_plot(axes=None,x=None,y=None):
    axes=check_axes(axes)
    axes.bar(left=x,height=y)

But Is there a better/clean/more pythonic way? I guess I can use a decorator but didn't figure it out. could anybody give some idea?

Thanks!!

Answer 1

Here's how to do it with a decorator:

import matplotlib.pyplot as plt    

def check_axes(plot_fn):
    def _check_axes_wrapped_plot_fn(axes=None, x=None, y=None):
        if not axes:
            fig = plt.figure()
            axes = fig.add_subplot(111)
            return plot_fn(axes, x, y)
        else:
            return plot_fn(axes, x, y)
    return _check_axes_wrapped_plot_fn

@check_axes
def line_plot(axes, x=None, y=None):
    axes.plot(x, y)

@check_axes
def bar_plot(axes, x=None, y=None):
    axes.bar(left=x, height=y)

How it works: the @check_axes syntax redefines the name of the decorated function, e.g. line_plot to be a new function that's created by the decorator, i.e. _check_axes_wrapped_plot_fn. This "wrapped" function handles the axes-checking logic and then calls the original plot function.

If you want check_axes to be able to decorate any plot function that takes an axes as its first argument, not just those that also take only x and y arguments, you can use Python's handy * syntax for arbitrary argument lists:

def check_axes(plot_fn):
    def _check_axes_wrapped_plot_fn(axes=None, *args):
        if not axes:
            fig = plt.figure()
            axes = fig.add_subplot(111)
            return plot_fn(axes, *args)  # pass all args after axes
        else:
            return plot_fn(axes, *args)  # pass all args after axes
    return _check_axes_wrapped_plot_fn  

Now, whether any of this is "better/cleaner/more Pythonic" is probably a matter of debate, and depends on the larger context.

By the way, in the spirit of being "more Pythonic," I reformatted your code to be closer to the PEP8 style guide. Notice the spaces after commas in parameter lists, the spaces around the = assignment operator (but not when = used for function keyword parameters), and saying not axes instead of axes == None.