CODEWITHSUNDEEP
pythongeneratordecorator
Am1rr3zA
Am1rr3zA

Reputation: 7411

python - decorate a generator

I have a function in 2 different version; it reads a large file (I make it simple here and read a very small excel file).

Version 1: Read the whole file and return the list of rows
Version 2: Read it line by line with help of a generator

I want to decorate the output for these 2 functions and add something to the end of each row based on different logic, that's why I think I need a different customized decorator. Yet I can't figure out how can I achieve it by the help of decorator? especially when I have yield instead of return.

Version1:

@dec
def readxls():
    fileBook = xlrd.open_workbook('../decorator.xls')
    sh = fileBook.sheet_by_name("Sheet1")
    out = []
    for row_index in xrange(1, sh.nrows):
        out.append(sh.row_values(row_index))
    return out

Version 2:

@dec2
def readxls():
    fileBook = xlrd.open_workbook('../decorator.xls')
    sh = fileBook.sheet_by_name("Sheet1")
    for row_index in xrange(1, sh.nrows):
        yield sh.row_values(row_index)

let's say the excel file is like:

Col1    Col2    Col3
Val11   Val12   Val13
Val21   Val22   Val23

I want to decorate the output to get the following result :

Col1    Col2    Col3   0  Col1Col2
Val11   Val12   Val13  1  Val11Val12
Val21   Val22   Val23  2  Val21Val22

In order to get something like this as output how should be my dec1 and dec2 function?

Upvotes: 1

Views: 951

Answers (1)

Copperfield
Copperfield

Reputation: 8510

decorator are suppose to work a taking the result of the function manipulate it and give the new result, so knowing what is the result is the key, in this case, for that example the result is [['val11', 'val12', 'val13'], ['val21', 'val22', 'val23']] for version 1 and a generator with the elements of that for the second. With this knowledge we can proceed to make a decorator, like for example

from functools import wraps

def decorator1(fun):
    @wraps(fun)
    def wrapper(*args,**kwds):
        result = fun(*args,**kwds)
        for i,x in enumerate(result,1):
            x.extend( (i, x[0]+x[1]) )
        return result
    return wrapper

def decorator2(fun):
    @wraps(fun)
    def wrapper(*args,**kwds):
        for i,x in enumerate(fun(*args,**kwds),1):
            x.extend( (i, x[0]+x[1]) )
            yield x
    return wrapper

(here I use wraps to help with maintaining some meta data of the decorate function (functionality wise is not needed) and as guide to write the example)

In the first decorator, as the result is the whole list, I just add the extra stuff to each element and return it, and in the second I add the extra stuff as they come along maintaining the generator structure

decorated with those, the result is now [['val11', 'val12', 'val13', 1, 'val11val12'], ['val21', 'val22', 'val23', 2, 'val21val22']]

As an aside note, because yours 2 function do the same thing, I would rather keep the generator and when I need a list call list(readxls()) and also I would add 2 extra variables to the function signature with default value those string to make the function more flexible.

Upvotes: 2

Related Questions