python每期一练1015

"""
交通信号灯切换的规律依次是：红灯过后是绿色，绿灯过后是黄灯，黄灯过后是红灯。
请编写一个函数，接收一个表示当前灯颜色的字符串，返回下一个信号灯颜色。
示例：
输入：“red”，输出：“green”
题目难度：简单
"""

def update_light(current: str) -> str:
    lights = {
        'red': 'green',
        'green': 'yellow',
        'yellow': 'red',
    }
    return lights[current]


if __name__ == '__main__':
    assert update_light("red") == "green"
    assert update_light("green") == "yellow" 
    assert update_light("yellow") == "red"