Какая функция нужна для нахождения простого числа в Python?

import math


def is_prime(number):
    # список простых чисел начинается с 2, всё остальное можно сразу отмести
    if number <= 1:
        return False
    number_sqrt = int(math.sqrt(number)) 
    divisors = range(2, (number_sqrt + 1))
    # Если число не простое, то в отрезке от 1 до квадратного корня числа, точно будут его делители.
    for element in divisors:
        if number % element == 0:
            return False
    return True


is_prime(0) # False
is_prime(1) # False
is_prime(2) # True
is_prime(3) # True
is_prime(4) # False