Initial type annotations (#39)

* Variable annotation syntax is only supported in Python 3.6 and greater

.travis.yml

@@ -3,18 +3,18 @@ os: linux
 dist: xenial   # required for Python >= 3.7
 language: python
 python:
- - "3.4"
- - "3.5"
  - "3.6"
  - "3.7"
  - "3.8"
  - "nightly"
 install:
+ - pip install mypy
  - pip install coverage
  - pip install pylint
 script:
+ - mypy primes.py binary_search.py --strict
  - python test.py
  - coverage run --include=primes.py,test.py,binary_search.py test.py && coverage report -m --fail-under=100 --rcfile=coveragerc
  - python speed_test.py --all
  - python speed_test.py --fermat 6
- - pylint primes.py || true
+ - pylint primes.py binary_search.py || true

binary_search.py

@@ -3,16 +3,19 @@
 """
 
 from bisect import bisect_left, bisect_right
+from typing import Optional, Sequence, TypeVar
 
-def binary_search(haystack, needle, low=0, high=None): # can't use a to specify default for high
+T = TypeVar('T')
+
+def binary_search(haystack: Sequence[T], needle: T, low: int = 0, high: Optional[int] = None) -> int: # can't use a to specify default for high
     """
     Same as bisect_left but returns -1 if x is not in lst
     """
     high = high if high is not None else len(haystack) # high defaults to len(a)
     pos = bisect_left(haystack, needle, low, high) # find insertion position
     return pos if pos != high and haystack[pos] == needle else -1 # don't walk off the end
 
-def list_up_to(lst, limit):
+def list_up_to(lst: Sequence[T], limit: int) -> Sequence[T]:
     """
     lst - sorted list
     Returns elements of lst <= limit

mypy.ini

@@ -0,0 +1,6 @@
+[mypy]
+warn_return_any = True
+warn_unused_configs = True
+
+[mypy-numpy.*]
+ignore_missing_imports = True

primes.py

@@ -2,6 +2,7 @@
 Several prime number functions
 """
 from math import log
+from typing import Any, Iterable, List, Optional, Tuple, Sequence, Set, Union
 
 try:
     from math import gcd
@@ -12,13 +13,13 @@
 
 from binary_search import binary_search, list_up_to
 
-class Primes(list):
+class Primes(List[int]):
     """
     List subclass that supports slicing and membership checking, automatically
     generating new primes when needed
     """
 
-    def __contains__(self, item):
+    def __contains__(self, item: Any) -> bool:
         """
         Check if a number is prime:
         >>> primes = Primes()
@@ -27,7 +28,7 @@ def __contains__(self, item):
         """
         return is_prime(item, self)
 
-    def __getitem__(self, key):
+    def __getitem__(self, key: Any) -> Any:
         """
         Used in slicing to get a single prime or a sequence of primes
         """
@@ -51,21 +52,21 @@ def __getitem__(self, key):
 
         return super().__getitem__(key)
 
-    def index(self, prime):
+    def index(self, prime: int, start: int = 0, stop: Optional[int] = None) -> int:
         """
         The index of the prime
         """
         if not self or self[-1] < prime:
             super().extend(primes_up_to(prime, self)[len(self):])
-        return super().index(prime)
+        return super().index(prime, start, stop or len(self))
 
-def _first_multiple_of(num, above):
+def _first_multiple_of(num: int, above: int) -> int:
     """
     Returns first multiple of num >= above
     """
     return above + ((num-(above%num)) % num)
 
-def sieve_of_eratosthenes(limit, primes=None):
+def sieve_of_eratosthenes(limit: int, primes: Optional[Sequence[int]] = None) -> Sequence[int]:
     """
     Implementation of Sieve of Eratosthenes
 
@@ -109,12 +110,14 @@ def sieve_of_eratosthenes(limit, primes=None):
         # Remember that 2 is prime, as it isn't referred to in the list
         primes = [2]
 
+    l_primes = list(primes)
+
     # Only go up to the position of the sqrt as all composites <= limit have
     # a factor <= limit, so all composites will have been marked by sqrt(limit)
     for num in range(offset, int(limit**0.5)+1, 2):
         # Hasn't been crossed yet, so it's prime
         if lst[(num-offset) // 2]:
-            primes.append(num)
+            l_primes.append(num)
 
             lst[(num**2 - offset)//2::num] = False
 
@@ -123,13 +126,13 @@ def sieve_of_eratosthenes(limit, primes=None):
     # Start at the position of the square root + 1, rounded up to be odd
     # If primes passed in were > sqrt(limit), start at the position of the
     # highest known prime + 1
-    start = (max(primes[-1], int(limit ** 0.5)) + 1) | 1
+    start = (max(l_primes[-1], int(limit ** 0.5)) + 1) | 1
 
-    primes.extend(num for num in range(start, limit+1, 2) if lst[(num-offset) // 2])
+    l_primes.extend(num for num in range(start, limit+1, 2) if lst[(num-offset) // 2])
 
-    return primes
+    return l_primes
 
-def _range24(start, stop, step=2):
+def _range24(start: int, stop: int, step: int = 2) -> Iterable[int]:
     """
     Like range(), but step is alternating between 2 and 4
     (or 4 and 2 if step is initially 4)
@@ -139,7 +142,7 @@ def _range24(start, stop, step=2):
         start += step
         step = 2+(step&2)
 
-def sieve_of_atkin(limit):
+def sieve_of_atkin(limit: int) -> Sequence[int]:
     """
     Implementation of Sieve of Atkin
 
@@ -189,22 +192,22 @@ def sieve_of_atkin(limit):
 
     return res + [num for num in range(_first_multiple_of(2, limit_sqrt)+1, limit+1, 2) if lst[num]]
 
-def primes_up_to(limit, primes=None):
+def primes_up_to(limit: int, primes: Optional[Sequence[int]] = None) -> Sequence[int]:
     """
     Returns primes up to (and including) limit
 
     Uses (hopefully) the faster sieving algorithm available
     """
     return sieve_of_eratosthenes(limit, primes)
 
-def _trial_division(num, primes):
+def _trial_division(num: int, primes: Sequence[int]) -> bool:
     """
     Simple trial division algorithm, check if num is prime by dividing
     it by known primes
     """
     return all(num%p != 0 for p in list_up_to(primes, int(num ** 0.5)))
 
-def _miller_rabin_2(num):
+def _miller_rabin_2(num: int) -> bool:
     """
     Single application of the Miller-Rabin primality test base-2
 
@@ -230,7 +233,7 @@ def _miller_rabin_2(num):
 
     return False
 
-def _jacobi_symbol(a, n):
+def _jacobi_symbol(a: int, n: int) -> int:
     """
     Calculate the Jacobi symbol (a/n)
     """
@@ -246,7 +249,7 @@ def _jacobi_symbol(a, n):
         elif n % 8 in [1, 7]:
             return 1
     elif a < 0:
-        return (-1)**((n-1)//2) * _jacobi_symbol(-1*a, n)
+        return int((-1)**((n-1)//2) * _jacobi_symbol(-1*a, n))
 
     if a % 2 == 0:
         return _jacobi_symbol(2, n) * _jacobi_symbol(a // 2, n)
@@ -258,7 +261,7 @@ def _jacobi_symbol(a, n):
         else:
             return _jacobi_symbol(n, a)
 
-def _D_chooser(num):
+def _D_chooser(num: int) -> int:
     """
     Choose a D value suitable for the Baillie-PSW test
     """
@@ -268,8 +271,7 @@ def _D_chooser(num):
         D *= -1
     return D
 
-def _U_V_subscript(k, n, U, V, P, Q, D):
-    k, n, U, V, P, Q, D = map(int, (k, n, U, V, P, Q, D))
+def _U_V_subscript(k: int, n: int, U: int, V: int, P: int, Q: int, D: int) -> Tuple[int, int]:
     digits = list(map(int, str(bin(k))[2:]))
     subscript = 1
     for digit in digits[1:]:
@@ -289,7 +291,7 @@ def _U_V_subscript(k, n, U, V, P, Q, D):
             U, V = U % n, V % n
     return U, V
 
-def _lucas_pp(num):
+def _lucas_pp(num: int) -> bool:
     """
     Perform the Lucas probable prime test
 
@@ -324,7 +326,7 @@ def _lucas_pp(num):
 
         return False
 
-def is_prime(num, primes=None):
+def is_prime(num: int, primes: Optional[Sequence[int]] = None) -> bool:
     """
     Returns True if num is a prime number, False if it is not
 
@@ -346,14 +348,13 @@ def is_prime(num, primes=None):
 
     return _miller_rabin_2(num) and _lucas_pp(num)
 
-def n_primes(num, primes=None):
+def n_primes(num: int, primes: Optional[Sequence[int]] = None) -> Sequence[int]:
     """
     Returns the first num primes
 
     Can pass in known primes to decrease execution time
     """
-    if not primes:
-        primes = Primes()
+    primes = primes or Primes()
 
     if len(primes) < num:
         if num < 6:
@@ -383,22 +384,22 @@ def n_primes(num, primes=None):
         primes = primes_up_to(upper_bound, primes)
     return primes[:num]
 
-def nth_prime(num, primes=None):
+def nth_prime(num: int, primes: Optional[Sequence[int]] = None) -> int:
     """
     Returns the numth prime (e.g. the 3rd prime, the 6th prime)
 
     Can pass in known primes to decrease execution time
     """
     return n_primes(num, primes)[-1]
 
-def composites_up_to(limit, primes=None):
+def composites_up_to(limit: int, primes: Optional[Sequence[int]] = None) -> Sequence[int]:
     """
     Returns all composite (non-prime greater than 1) numbers up to (and including) limit
 
     Can pass in known primes to decrease execution time
     """
     primes = primes_up_to(limit, primes)
-    composites = []
+    composites: List[int] = []
     for prime1, prime2 in zip(primes, primes[1:]):
         # Add numbers between primes to composites
         composites.extend(range(prime1+1, prime2))
@@ -407,7 +408,7 @@ def composites_up_to(limit, primes=None):
         composites.extend(range(primes[-1]+1, limit+1))
     return composites
 
-def next_prime(primes):
+def next_prime(primes: List[int]) -> int:
     """
     Given primes, returns the next prime
 
@@ -422,71 +423,67 @@ def next_prime(primes):
     for num in range(_first_multiple_of(2, primes[-1])+1, 2*primes[-1], 2):
         if is_prime(num, primes):
             return num
+
+    raise RuntimeError("Unreachable code") # pragma: no cover
 
-def primes_with_difference_up_to(limit, difference, primes=None):
+def primes_with_difference_up_to(limit: int, difference: int, primes: Optional[Sequence[int]] = None) -> Iterable[Tuple[int, int]]:
     """
     Primes with difference up to limit
     """
-    if not primes:
-        primes = Primes()
+    primes = primes or Primes()
 
     return ((prime, prime+difference) for prime in primes_up_to(limit-difference, primes)
             if is_prime(prime+difference, primes))
 
-def twin_primes_up_to(limit, primes=None):
+def twin_primes_up_to(limit: int, primes: Optional[Sequence[int]] = None) -> Iterable[Tuple[int, int]]:
     """
     Primes with difference 2 up to limit
     """
     return primes_with_difference_up_to(limit, 2, primes)
 
-def cousin_primes_up_to(limit, primes=None):
+def cousin_primes_up_to(limit: int, primes: Optional[Sequence[int]] = None) -> Iterable[Tuple[int, int]]:
     """
     Primes with difference 4 up to limit
     """
     return primes_with_difference_up_to(limit, 4, primes)
 
-def sexy_primes_up_to(limit, primes=None):
+def sexy_primes_up_to(limit: int, primes: Optional[Sequence[int]] = None) -> Iterable[Tuple[int, int]]:
     """
     Primes with difference 6 up to limit
     """
     return primes_with_difference_up_to(limit, 6, primes)
 
-def prime_triplets_up_to(limit, primes=None):
+def prime_triplets_up_to(limit: int, primes: Optional[Sequence[int]] = None) -> Iterable[Tuple[int, int, int]]:
     """
     Prime triplets up to limit
     """
-    if not primes:
-        primes = Primes()
+    primes = primes or Primes()
 
     for prime in primes_up_to(limit-6, primes):
         if is_prime(prime+2, primes) and is_prime(prime+6, primes):
             yield (prime, prime+2, prime+6)
         if is_prime(prime+4, primes) and is_prime(prime+6, primes):
             yield (prime, prime+4, prime+6)
 
-def prime_quadruplets_up_to(limit, primes=None):
+def prime_quadruplets_up_to(limit: int, primes: Optional[Sequence[int]] = None) -> Iterable[Tuple[int, int, int, int]]:
     """
     Prime quadruplets up to limit
     """
-    if not primes:
-        primes = Primes()
+    primes = primes or Primes()
 
     for prime in primes_up_to(limit-8, primes):
         if is_prime(prime+2, primes) and is_prime(prime+6, primes) and is_prime(prime+8, primes):
             yield (prime, prime+2, prime+6, prime+8)
 
-def prime_gaps_up_to(limit, primes=None):
+def prime_gaps_up_to(limit: int, primes: Optional[Sequence[int]] = None) -> Iterable[int]:
     """
     Difference between successive primes up to limit
     """
-    if not primes:
-        primes = Primes()
-
     primes = primes_up_to(limit, primes)
     for prime1, prime2 in zip(primes, primes[1:]):
         yield prime2 - prime1
 
-def brents_rho(num, starting_point=2):
+def brents_rho(num: int, starting_point: int = 2) -> int:
     """
     Return a factor of num using Brent's variant of Pollard's rho algorithm, or num if one is not found
     """
@@ -519,15 +516,14 @@ def brents_rho(num, starting_point=2):
 
     return g
 
-def factorise(num, include_trivial=False, primes=None):
+def factorise(num: int, include_trivial: bool = False, primes: Optional[Sequence[int]] = None) -> Set[int]:
     """
     Factorise a number
 
     Returns the prime factors of num
     Excludes trivial factors unless include_trivial = True
     """
-    if not primes:
-        primes = Primes()
+    primes = primes or Primes()
 
     factors = set([1, num]) if include_trivial else set()