python 函数 判断是否为素数

def su(n):
    i = 2
    flag = True
    while i <= n - 1:
        if n % i == 0:
            flag = False
        i += n
    if flag:
        return True
    else:
        return False
su = su(int(input("请输入一个数：")))
if su:
    print("是素数")
else:
    print("不是素数")