Search code examples
pythonpython-3.xgeneratorgenerator-expression

Assign two variables via generator expression w/ conditional?


The below code has a dictionary containing various cheeses, as well as their quantities. Based off a pre-determined list of sale items, the code then prints the total quantity of cheeses which are on sale versus full price.

I'm using generator expressions to calculate the totals, but I'm wondering if there's a way to condense my code to assign both the sale_count and full_price_count variables at once with some sort of if-else condition, as the code for the generators is practically the same.

cheeses = {'gouda': 3, 'cheddar': 7, 'american': 2, 'mozzarella': 5}
on_sale = ['american', 'blue cheese', 'cheddar', 'provolone', 'swiss']

# if the cheese is on sale, add its quantity to sale_count
# otherwise, add its quantity to full_price_count
sale_count = sum(qty for (cheese, qty) in cheeses.items() if cheese in on_sale)
full_price_count = sum(qty for (cheese, qty) in cheeses.items() if cheese not in on_sale)

print("Sale count: {}\nFull price count: {}".format(sale_count, full_price_count))

Solution

  • It could be done in a single expression as:

    functools.reduce(
        lambda x, y: (x[0] + y[0], x[1] + y[1]),
        ((qty, 0) if cheese in on_sale else (0, qty) for cheese, qty in cheeses.items()), 
        (0, 0))
    

    But then this, as with other potential answers, might really answer why things don't always have to be simplified into one expression when two is perfectly clear.