viết chương trình kiểm tra số nguyên tố python nha

def prime(n):
    ok = None
    if n < 2 : ok = False
    else : ok = True 
    for i in range(2,n):
        if n % i == 0 : 
            ok = False  
            break 
    print(ok)
    return  
n=int(input())
prime(n)