Search code examples
pythonfunctionpython-decorators

Python - can I use a decorator for this?

def f(a,b):
    lista = []
    listb = []
    lista.append(a**2)
    listb.append(b**2)
    return lista, listb

but I also want to return an array made from the two lists.

so like:

def f2(a,b):
    lista = []
    listb = []
    lista.append(a**2)
    listb.append(b**2)
    tr = np.array([np.array(lista), np.array(listb)]).T
    return tr

Both functions share the same arguments, but I would not like to pile up both the lists and the array as returned values.

Can I somehow use a decorator no f so as to define another function that just takes the argument of it and does another operation?

Solution

  • Yes, you can use the following decorator:

    def tr(func):
    def wrapper(a, b):
        lista, listb = func(a, b)
        return np.array([np.array(lista), np.array(listb)]).T
    return wrapper

    so that f2 can be defined as simply:

    @tr
def f2(a, b):
    return f(a, b)