CODEWITHSUNDEEP
pythonlist-comprehension
Kartik
Kartik

Reputation: 8703

List comprehension based on choices

Basically, if I were to write a function with variable return elements, like so:

def func(elem1=True, elem2=True, elem3=True, elem4=False):
    x = MyClass()
    ret = []
    if elem1:
        ret.extend([x.func1()])
    if elem2:
        ret.extend([x.obj1])
    if elem3:
        ret.extend([x.func2().attr1])
    if elem4:
        ret.extend(x.list_obj3)
    return ret

Things get rather long and windy. Is it possible to do something like this perhaps:

def func(elem1=True, elem2=True, elem3=True, elem4=False):
    x = MyClass()
    return [x.func1() if elem1,
            x.obj1 if elem2,
            x.func2().attr1 if elem3,
            x.list_obj3 if elem4]

How neat is that!?

I know this can be done:

def func(elem1=True, elem2=True, elem3=True, elem4=False):
    x = MyClass()
    ret = [x.func1(), x.obj1, x.func2().attr1, x.list_obj3]
    choices = [elem1, elem2, elem3, elem4]
    return [r for i, r in enumerate(ret) if choices[i]]

but I would like to not calculate the elements if the user does not want them; it is a little expensive to calculate some of them.

Upvotes: 2

Views: 91

Answers (2)

Neapolitan
Neapolitan

Reputation: 2163

Asking a slightly different question, can you get behaviour like matlab/octave, where you only calculate the first two results if you are assigning to two variables, without computing results 3 and 4?

For example:

a, b = func()

Python can't quite do it since func() doesn't know how many return values it wants, but you can get close using:

from itertools import islice
def func():
    x = MyClass()
    yield x.fun c1()
    yield x.obj1
    yield x.func2().attr1
    yield x.list_obj3

a, b = islice(func(), 2)

I'm not sure it is better, but you could add array indexing semantics using a decorator, which would allow you to write:

@sliceable
def func():
    ...
a, b = func()[:2]

This is easy enough to implement:

from itertools import islice
class SlicedIterator(object):
    def __init__(self, it):
        self.it = it
    def __iter__(self): 
        return self.it
    def __getitem__(self, idx):
        if not isinstance(idx, slice):
            for _ in range(idx): next(self.it)
            return next(self.it)
        return list(islice(self.it, idx.start, idx.stop, idx.step))
def sliceable(f):
    def wraps(*args, **kw):
        return SlicedIterator(f(*args, **kw))
    return wraps

Testing:

@sliceable
def f():
    print("compute 1")
    yield 1
    print("compute 2")
    yield 2
    print("compute 3")
    yield 3
    print("compute 4")
    yield 4

print("== compute all four")
a, b, c, d = f()
print("== compute first two")
a, b = f()[:2]
print("== compute one only")
a = f()[0]
print("== all as a list")
a = f()[:]

gives:

== compute all four
compute 1
compute 2
compute 3
compute 4
== compute first two
compute 1
compute 2
== compute one only
compute 1
== all as a list
compute 1
compute 2
compute 3
compute 4

Upvotes: 1

Neapolitan
Neapolitan

Reputation: 2163

If you hide your operations in lambdas then you can use lazy evaluation:

def func(elem1=True, elem2=True, elem3=True, elem4=False):
    x = MyClass()
    return [L() for inc,L in (
            (elem1, lambda: x.func1()),
            (elem2, lambda: x.obj1),
            (elem3, lambda: x.func2().attr1),
            (elem4, lambda: x.list_obj3),
            ) if inc]

Upvotes: 5

Related Questions