New exercise: python nth-prime

EXERCISES.txt

@@ -21,6 +21,7 @@ gigasecond
 triangle
 scrabble-score
 sieve
+nth-prime
 luhn
 roman-numerals
 binary

nth-prime/example.py

@@ -0,0 +1,31 @@
+from itertools import count
+from math import sqrt
+
+
+def nth_prime(n):
+    known = []
+    candidates = prime_candidates()
+
+    def is_prime(m):
+        sqrt_m = sqrt(m)
+        for k in known:
+            if k > sqrt_m:
+                return True
+            elif m % k == 0:
+                return False
+        return True
+
+    while len(known) < n:
+        x = next(candidates)
+        if is_prime(x):
+            known.append(x)
+
+    return known[n - 1]
+
+
+def prime_candidates():
+    yield 2
+    yield 3
+    for n in count(6, 6):
+        yield n - 1
+        yield n + 1

nth-prime/nth_prime_test.py

@@ -0,0 +1,25 @@
+try:
+    from prime import nth_prime
+except ImportError:
+    raise SystemExit('Could not find prime.py. Does it exist?')
+
+import unittest
+
+
+class NthPrimeTests(unittest.TestCase):
+    def test_first_prime(self):
+        self.assertEqual(2, nth_prime(1))
+
+    def test_sixth_prime(self):
+        self.assertEqual(13, nth_prime(6))
+
+    def test_first_twenty_primes(self):
+        self.assertEqual([2,3,5,7,11,13,17,19,23,29,31,37,41,43,47,53,59,61,67,71],
+                         [nth_prime(n) for n in range(1, 21)])
+
+    def test_prime_no_10000(self):
+        self.assertEqual(104729, nth_prime(10000))
+
+
+if __name__ == '__main__':
+    unittest.main()