Prime field arithmetic in Python

fp.py

#!/usr/bin/env python
#
# Author: Madars Virza <madars@mit.edu>
# License: MIT
#

import random
import operator

def ext_gcd(a, b):
    """Run extended Euclidean algorithm for inputs a and b and return
    triple (d, x, y) such that a*x+b*y = d = gcd(a,b)."""
    if a == 0:
        return (b, 0, 1)
    if b == 0:
        return (a, 1, 0)
    q, r = divmod(b, a)
    d, xp, yp = ext_gcd(r, a) # in our later calls a < b
    # (b % a)*x' + a*y' = d
    # (b - q*a)*x' + a*y' = d
    # a*(y'-q*x') + b*x' = d
    # and, of course, gcd(a, b) = gcd(b % a, a)
    return (d, yp-q*xp, xp)

def invmod(a, n):
    """Return modular the inverse of an input a modulo n for n not necessarily
    prime."""
    d, x, y = ext_gcd(a, n)
    if abs(d) != 1:
        raise ValueError("modulo must be coprime with the element")
    return x * d

class Fp_model(object):
    """Implementation of arithmetic modulo p, a prime."""
    def __init__(self, val, modulus):
        self.modulus = modulus
        self.val = val % modulus

    def __repr__(self):
        return "%d" % self.val

    def __neg__(self):
        return Fp_model(-self.val, self.modulus)

    def inverse(self):
        return Fp_model(invmod(self.val, self.modulus), self.modulus)

    def _op(self, other, op):
        if not isinstance(other, Fp_model) or self.modulus != other.modulus:
            raise ValueError("You can only do this between two Fp elements of matching modulus.")
        return Fp_model(op(self.val, other.val), self.modulus)

    __add__ = lambda self, other: self._op(other, operator.add)
    __sub__ = lambda self, other: self._op(other, operator.sub)
    __mul__ = lambda self, other: self._op(other, operator.mul)
    __div__ = lambda self, other: self._op(other.inverse(), operator.mul)
    __truediv__ = __div__

    def __pow__(self, power):
        if not isinstance(power, int) and not isinstance(power, long):
            raise ValueError("Power must be an integer.")
        return Fp_model(pow(self.val, power, self.modulus), self.modulus)

    def __eq__(self, other):
        return isinstance(other, Fp_model) and self.modulus == other.modulus and self.val == other.val

    def __ne__(self, other):
        return not (self == other)

def Fp_factory(p):
    def Fp(val):
        return Fp_model(val, p)
    return Fp

def test():
    # Fr for default curve choice in libsnark
    bn128_r = 21888242871839275222246405745257275088548364400416034343698204186575808495617
    Fr = Fp_factory(bn128_r)

    one, two, three, four, five, six = Fr(1), Fr(2), Fr(3), Fr(4), Fr(5), Fr(6)

    assert one + two == three
    assert two + two == four
    assert two * three == six
    assert two * three != five
    assert six * two.inverse() == three
    assert (-two) * (two.inverse()) == -one
    assert two**bn128_r == two
    assert two**(bn128_r-1) == one
    assert two**(bn128_r-2) == two.inverse()

    print "two:", two
    print "two.inverse():", two.inverse()

if __name__ == '__main__':
    test()