素数を判定するpythonコード

main.py

"""
Input
[1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20]

Prime number
[2,3,5,7,11,13,17,19]

Non Prime number
[1,4,6,8,9,12,14,15,16,18,20]

"""

# 素数を判定する関数(1以下の値は素数ではないのでfalseを返す)


def is_prime_v1(num: int) -> bool:
    if num <= 1:
        return False

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

    return True


# 処理数を減らして速度を上げた素数判定関数
def is_prime_v2(num: int) -> bool:
    if num <= 1:
        return False

    for i in range(2, math.floor(math.sqrt(num) + 1)):
        if num % i == 0:
            return False

    return True
    # >>> print([i for i in range(2, math.floor(math.sqrt(37) + 1))])
    #[2, 3, 4, 5, 6]


# 本ファイル実行時に出力
if __name__ == '__main__':
    # print(is_prime_v1(10))
    import math
    import time
    import random
    numbers = [random.randint(0, 1000) for _ in range(10000)]

    start = time.time()
    for num in numbers:
        is_prime_v1(num)
    print('v1', time.time() - start)

    start = time.time()
    for num in numbers:
        is_prime_v2(num)
    print('v2', time.time() - start)

unittestするpythonコード

unittest.py


import unittest

# 多数の関数を作るので別名付与
from main import is_prime_v1 as is_prime


class primeTest(unittest.TestCase):
    # Prime numberでtrueを返さなかったら失敗
    def test_is_prime_ok(self):
        for i in [2,3,5,7,11,13,17,19]:
            self.assertTrue(is_prime(i))


    # Non Prime numberでFalseを返さなかったら失敗
    def test_is_prime_no(self):
        for i in [1,4,6,8,9,12,14,15,16,18,20]:
            self.assertTrue(is_prime(i))


    # マイナス値でFalseを返さなかったら失敗
    def test_is_prime_negative(self):
            self.assertTrue(is_prime(-1))


    # typeerrorをraise
    def test_is_prime_typeerror(self):
            with self.assertRaise(typeError):
                is_prime('string')


# 本ファイル実行時に出力
if __name__ == '__main__':
    unittest.main()