Search code examples
pythonfor-loopdry

How to write DRY Python For Loop


I have a cannabis dataset that has a column for "Effects" and I'm trying to add a binary "nice_buds" column for strains that do not include certain effects. This is the code:

nice_buds = []
undesired_effects = ["Sleepy", "Hungry", "Giggly", "Tingly", "Aroused", "Talkative"]

for row in sample["Effects"]:
    if "Sleepy" not in row and "Hungry" not in row and "Giggly" not in row and "Tingly" not in row and "Aroused" not in row and "Talkative" not in row:
        nice_buds.append(1)
    else:
        nice_buds.append(0)

sample["nice_buds"] = nice_buds

As of now, the undesired_effects list is doing nothing, and the code works perfectly fine in terms of giving me the desired output.

My question though is if there is a more "Pythonic" or "DRY" way to go about this ...


Solution

  • You could use all() with a generator expression to simplify the if-statement

    nice_buds = []
    undesired_effects = ["Sleepy", "Hungry", "Giggly", "Tingly", "Aroused", "Talkative"]
    
    for row in sample["Effects"]:
        if all(effect not in row for effect in undesired_effects):
            nice_buds.append(1)
        else:
            nice_buds.append(0)
    
    sample["nice_buds"] = nice_buds
    

    Or use any() & check for the presence of an effect:

    nice_buds = []
    undesired_effects = ["Sleepy", "Hungry", "Giggly", "Tingly", "Aroused", "Talkative"]
    
    for row in sample["Effects"]:
        if any(effect in row for effect in undesired_effects):
            nice_buds.append(0)
        else:
            nice_buds.append(1)
    
    sample["nice_buds"] = nice_buds