Raspberry Pi Traffic Lights

import gpiozero as io #shortened version, but retaining full class path
import time

WAIT = 4 # Constant for time to keep light on red or green

def off(leds):
    '''Pass dictionary of LED defined in main and set all to off'''
    for led in leds.values():
        led.off()

def on(leds, colours):
    '''Pass dictionary of LED defined in main and list of colours to switch on'''
    for colour in colours:
        leds[colour].on()
    time.sleep(1.5) # built in timer

def main():
    # define dictionary of LEDs
    leds = {'red':io.LED(17), 'yellow':io.LED(27), 'green':io.LED(22)}
    while True:
        off(leds)
        on(leds, ['red'])
        time.sleep(WAIT)
        on(leds, ['red', 'yellow'])
        off(leds)
        on(leds, ['green'])
        time.sleep(WAIT)
        off(leds)
        on(leds, ['yellow'])

main()