定义一个函数判断一个数是否为素数 python

def is_prime(num):
    if num == 1:
        return False
    elif num == 2:
        return True

    for i in range(2, num):
        if num % i == 0:
            return False

    return True


num = int(input('输入一个数 '))

if is_prime(num):
    print('是素数')
else:
    print('不是素数')