Pswiss87
Pswiss87

Reputation: 755

How to use python generator expressions to create a oneliner to run a function multiple times and get a list output

I am wondering if there is there is a simple Pythonic way (maybe using generators) to run a function over each item in a list and result in a list of returns?

Example:

def square_it(x):
    return x*x

x_set = [0,1,2,3,4]
squared_set = square_it(x for x in x_set)

I notice that when I do a line by line debug on this, the object that gets passed into the function is a generator.

Because of this, I get an error: TypeError: unsupported operand type(s) for *: 'generator' and 'generator'

I understand that this generator expression created a generator to be passed into the function, but I am wondering if there is a cool way to accomplish running the function multiple times only by specifying an iterable as the argument? (without modifying the function to expect an iterable).

It seems to me that this ability would be really useful to cut down on lines of code because you would not need to create a loop to fun the function and a variable to save the output in a list.

Thanks!

Upvotes: 3

Views: 2941

Answers (5)

Ramchandra Apte
Ramchandra Apte

Reputation: 4079

There's a builtin function, map(), for this common problem.

>>> map(square_it, x_set)
[0,1,4,9,16] # On Python 3, a generator is returned.

Alternatively, one can use a generator expression, which is memory-efficient but lazy (meaning the values will not be computed now, only when needed):

>>> (square_it(x) for x in x_set)
    <generator object <genexpr> at ...>

Similarly, one can also use a list comprehension, which computes all the values upon creation, returning a list.

Additionally, here's a comparison of generator expressions and list comprehensions.

Upvotes: 4

Mauricio Abreu
Mauricio Abreu

Reputation: 155

Note that there is a difference between list comprehension returning a list

squared_set = [square_it(x) for x in x_set]

and returning a generator that you can iterate over it:

squared_set = (square_it(x) for x in x_set)

Upvotes: 1

Matt Anderson
Matt Anderson

Reputation: 19779

As the other answers have suggested, I think it is best (most "pythonic") to call your function explicitly on each element, using a list or generator comprehension.

To actually answer the question though, you can wrap your function that operates over scalers with a function that sniffs the input, and has different behavior depending on what it sees. For example:

>>> import types
>>> def scaler_over_generator(f):
...     def wrapper(x):
...         if isinstance(x, types.GeneratorType):
...             return [f(i) for i in x]
...         return f(x)
...     return wrapper
>>> def square_it(x):
...     return x * x
>>> square_it_maybe_over = scaler_over_generator(square_it)
>>> square_it_maybe_over(10)
    100
>>> square_it_maybe_over(x for x in range(5))
    [0, 1, 4, 9, 16]

I wouldn't use this idiom in my code, but it is possible to do.

You could also code it up with a decorator, like so:

>>> @scaler_over_generator
... def square_it(x):
...     return x * x
>>> square_it(x for x in range(5))
    [0, 1, 4, 9, 16]

If you didn't want/need a handle to the original function.

Upvotes: 1

Matt Bryant
Matt Bryant

Reputation: 4961

You want to call the square_it function inside the generator, not on the generator.

squared_set = (square_it(x) for x in x_set)

Upvotes: 2

BrenBarn
BrenBarn

Reputation: 251428

You want a list comprehension:

squared_set = [square_it(x) for x in x_set]

Upvotes: 5

Related Questions