String Validator

• User can input 'c' or 'f' to choose unit of measure, so here I need a function to assert if the input is valid. If the input is not equal to 'c' or 'f', user need to re-input, else, return the input unit.

def input_unit():
unit = ''
while True:
    unit = input('请输入温度单位>: ')
    unit = unit.strip().lower()
    if unit not in ['c', 'f']:
        print('输入错误，请重新输入!')
    else:
        break

return unit