Double for loop to Python list comprehension

Question

I have the following (simplified structure):

class Wafer:
    def __init__(self, df):
        self.dose_fields = df

class Data:
    def __init__(self, add_zero):
        self.wafers = []
        self.wafers.append(Wafer([1.0, 1.5]))
        self.wafers.append(Wafer([2.0, 2.0]))
        self.wafers.append(Wafer([3.0, 3.5]))
        self.wafers.append(Wafer([4.0, 4.5]))

        if add_zero:
            self.wafers.append(Wafer([5.0, 0.0]))

I want to have a function that returns False when there is a zero in one of the dose_fields values.

like this:

data = Data(True)        
res = True
for wafer in data.wafers:
    for fld in wafer.dose_fields:
        if fld == 0.0:
            res = False
print res

data = Data(False)        
res = True
for wafer in data.wafers:
    for fld in wafer.dose_fields:
        if fld == 0.0:
            res = False
print res

This prints False resp. True (what I want).

However, when I try to use list comprehension:

data = Data(True)
print [df for df in wafer.dose_fields for wafer in data.wafers]
print res

data = Data(False)
print [df for df in wafer.dose_fields for wafer in data.wafers]
print res

I get the result

[4.0, 4.0, 4.0, 4.0, 4.0, 4.5, 4.5, 4.5, 4.5, 4.5]
True

resp

[5.0, 5.0, 5.0, 5.0, 5.0, 0.0, 0.0, 0.0, 0.0, 0.0]
True

Somehow I get repeats of the last item, while I want the complete list and check afterwards if there is a 0.0 inside.

What should I do to 'repair' the list comprehension statement?

Answer 1

Nested list comprehensions can be a bit confusing. You actually need to write them in the same order as the for loops they replace: so it should be

[df for wafer in data.wafers for df in wafer.dose_fields]

Answer 2

If testing for the presence of 0.0 is really all you want, you should use the in operator. Not only is it far more readable, it's much faster. You can then use any , too

Instead of:

res = True
for wafer in data.wafers:
    for fld in wafer.dose_fields:
        if fld == 0.0:
            res = False

do

res= not any([0.0 in wafer.dose_fields for wafer in data.wafers])