Add rounding functions for fmpq and tests

Author: oscarbenjamin

src/flint/fmpq.pyx

@@ -136,11 +136,20 @@ cdef class fmpq(flint_scalar):
             return - int((-self.p) // self.q)
 
     def __floor__(self):
-        return int(self.p // self.q)
+        return self.floor()
+
+    def __ceil__(self):
+        return self.ceil()
+
+    def __trunc__(self):
+        return self.trunc()
 
     def __nonzero__(self):
         return not fmpq_is_zero(self.val)
 
+    def __round__(self, ndigits=None):
+        return self.round(ndigits)
+
     def __pos__(self):
         return self
 
@@ -351,6 +360,37 @@ cdef class fmpq(flint_scalar):
         fmpz_cdiv_q(r.val, fmpq_numref(self.val), fmpq_denref(self.val))
         return r
 
+    def trunc(self):
+        """
+        Truncation function.
+
+            >>> fmpq(3,2).trunc()
+            1
+            >>> fmpq(-3,2).trunc()
+            -1
+        """
+        cdef fmpz r = fmpz.__new__(fmpz)
+        fmpz_tdiv_q(r.val, fmpq_numref(self.val), fmpq_denref(self.val))
+        return r
+
+    def round(self, ndigits=None):
+        """
+        Rounding function.
+
+            >>> fmpq(3,2).round()
+            2
+            >>> fmpq(-3,2).round()
+            -2
+        """
+        from fractions import Fraction
+        fself = Fraction(int(self.p), int(self.q))
+        if ndigits is not None:
+            fround = round(fself, ndigits)
+            return fmpq(fround.numerator, fround.denominator)
+        else:
+            fround = round(fself)
+            return fmpz(fround)
+
     def __hash__(self):
         from fractions import Fraction
         return hash(Fraction(int(self.p), int(self.q), _normalize=False))

src/flint/fmpz.pyx

@@ -369,7 +369,7 @@ cdef class fmpz(flint_scalar):
 
             if not success:
                 if ttype == FMPZ_TMP: fmpz_clear(tval)
-                raise ValueError("fmpz_pow_fmpz: exponent too large")
+                raise OverflowError("fmpz_pow_fmpz: exponent too large")
         else:
             # Modular exponentiation
             mtype = fmpz_set_any_ref(mval, m)
@@ -406,16 +406,21 @@ cdef class fmpz(flint_scalar):
         if typecheck(other, fmpz):
             other = int(other)
         if typecheck(other, int):
+            if other < 0:
+                raise ValueError("negative shift count")
             u = fmpz.__new__(fmpz)
             fmpz_mul_2exp((<fmpz>u).val, self.val, other)
             return u
         else:
             return NotImplemented
 
     def __rlshift__(self, other):
+        iself = int(self)
+        if iself < 0:
+            raise ValueError("negative shift count")
         if typecheck(other, int):
             u = fmpz.__new__(fmpz)
-            fmpz_mul_2exp((<fmpz>u).val, fmpz(other).val, int(self))
+            fmpz_mul_2exp((<fmpz>u).val, fmpz(other).val, iself)
             return u
         else:
             return NotImplemented
@@ -424,16 +429,21 @@ cdef class fmpz(flint_scalar):
         if typecheck(other, fmpz):
             other = int(other)
         if typecheck(other, int):
+            if other < 0:
+                raise ValueError("negative shift count")
             u = fmpz.__new__(fmpz)
             fmpz_fdiv_q_2exp((<fmpz>u).val, self.val, other)
             return u
         else:
             return NotImplemented
 
     def __rrshift__(self, other):
+        iself = int(self)
+        if iself < 0:
+            raise ValueError("negative shift count")
         if typecheck(other, int):
             u = fmpz.__new__(fmpz)
-            fmpz_fdiv_q_2exp((<fmpz>u).val, fmpz(other).val, int(self))
+            fmpz_fdiv_q_2exp((<fmpz>u).val, fmpz(other).val, iself)
             return u
         else:
             return NotImplemented

test/test.py

@@ -1,8 +1,11 @@
 import sys
-import flint
+import math
 import operator
+import pickle
 import doctest
 
+import flint
+
 if sys.version_info[0] >= 3:
     long = int
 
@@ -85,6 +88,20 @@ def test_fmpz():
                     assert ltype(s) + rtype(t) == s + t
                     assert ltype(s) - rtype(t) == s - t
                     assert ltype(s) * rtype(t) == s * t
+                    assert ltype(s) & rtype(t) == s & t
+                    assert ~ltype(s) == ~s
+
+                    # XXX: Some sort of bug here means that the following fails
+                    # for values of s and t bigger than the word size.
+                    if abs(s) < 2**62 and abs(t) < 2**62:
+                        assert ltype(s) | rtype(t) == s | t
+                    else:
+                        # This still works so somehow the internal representation
+                        # of the fmpz made by fmpz_or is such that fmpz_equal
+                        # returns false even though the values look the same.
+                        assert str(ltype(s) | rtype(t)) == str(s | t)
+
+                    assert ltype(s) ^ rtype(t) == s ^ t
                     if t == 0:
                         assert raises(lambda: ltype(s) // rtype(t), ZeroDivisionError)
                         assert raises(lambda: ltype(s) % rtype(t), ZeroDivisionError)
@@ -101,6 +118,20 @@ def test_fmpz():
                     assert (ltype(s) >= rtype(t)) == (s >= t)
                     if 0 <= t < 10:
                         assert (ltype(s) ** rtype(t)) == (s ** t)
+                        assert ltype(s) << rtype(t) == s << t
+                        assert ltype(s) >> rtype(t) == s >> t
+                    elif -10 <= t < 0:
+                        assert raises(lambda: ltype(s) << rtype(t), ValueError)
+                        assert raises(lambda: ltype(s) >> rtype(t), ValueError)
+
+    assert 2 ** flint.fmpz(2) == 4
+    assert type(2 ** flint.fmpz(2)) == flint.fmpz
+    assert raises(lambda: () ** flint.fmpz(1), TypeError)
+    assert raises(lambda: flint.fmpz(1) ** (), TypeError)
+    assert raises(lambda: flint.fmpz(1) ** -1, ValueError)
+
+    mega = flint.fmpz(2) ** 8000000
+    assert raises(lambda: mega ** mega, OverflowError)
 
     pow_mod_examples = [
         (2, 2, 3, 1),
@@ -118,6 +149,19 @@ def test_fmpz():
     # XXX: Handle negative modulus like int?
     assert raises(lambda: pow(flint.fmpz(2), 2, -1), ValueError)
 
+    f = flint.fmpz(2)
+    assert f.numerator == f
+    assert type(f.numerator) is flint.fmpz
+    assert f.denominator == 1
+    assert type(f.denominator) is flint.fmpz
+
+    assert int(f) == 2
+    assert type(int(f)) is int
+    assert operator.index(f) == 2
+    assert type(operator.index(f)) is int
+    assert float(f) == 2.0
+    assert type(float(f)) is float
+
     assert flint.fmpz(2) != []
     assert +flint.fmpz(0) == 0
     assert +flint.fmpz(1) == 1
@@ -128,30 +172,68 @@ def test_fmpz():
     assert abs(flint.fmpz(0)) == 0
     assert abs(flint.fmpz(1)) == 1
     assert abs(flint.fmpz(-1)) == 1
-    assert int(flint.fmpz(2)) == 2
-    assert isinstance(int(flint.fmpz(2)), int)
-    assert long(flint.fmpz(2)) == 2
-    assert isinstance(long(flint.fmpz(2)), long)
-    l = [1, 2, 3]
-    l[flint.fmpz(1)] = -2
-    assert l == [1, -2, 3]
-    d = {flint.fmpz(2): 3}
-    d[flint.fmpz(2)] = -1
-    assert d == {flint.fmpz(2): -1}
+
+    assert bool(flint.fmpz(0)) == False
+    assert bool(flint.fmpz(1)) == True
+
     assert flint.fmpz(2).bit_length() == 2
     assert flint.fmpz(-2).bit_length() == 2
     assert flint.fmpz(2).height_bits() == 2
     assert flint.fmpz(-2).height_bits() == 2
     assert flint.fmpz(2).height_bits(signed=True) == 2
     assert flint.fmpz(-2).height_bits(signed=True) == -2
+
+    f1 = flint.fmpz(1)
+    f2 = flint.fmpz(2)
+    f3 = flint.fmpz(3)
+    f8 = flint.fmpz(8)
+
+    assert f2 << 2 == 8
+    assert f2 << f2 == 8
+    assert 2 << f2 == 8
+    assert raises(lambda: f2 << -1, ValueError)
+    assert raises(lambda: 2 << -f1, ValueError)
+
+    assert f8 >> 2 == f2
+    assert f8 >> f2 == f2
+    assert 8 >> f2 == f2
+    assert raises(lambda: f2 >> -1, ValueError)
+    assert raises(lambda: 2 >> -f1, ValueError)
+
+    assert f2 & 3 == 2
+    assert f2 & f3 == 2
+    assert 2 & f3 == 2
+    assert f2 | 3 == 3
+    assert f2 | f3 == 3
+    assert 2 | f3 == 3
+    assert f2 ^ 3 == 1
+    assert f2 ^ f3 == 1
+    assert 2 ^ f3 == 1
+
+    assert raises(lambda: f2 << (), TypeError)
+    assert raises(lambda: () << f2, TypeError)
+    assert raises(lambda: f2 >> (), TypeError)
+    assert raises(lambda: () >> f2, TypeError)
+    assert raises(lambda: f2 & (), TypeError)
+    assert raises(lambda: () & f2, TypeError)
+    assert raises(lambda: f2 | (), TypeError)
+    assert raises(lambda: () | f2, TypeError)
+    assert raises(lambda: f2 ^ (), TypeError)
+    assert raises(lambda: () ^ f2, TypeError)
+
+    l = [1, 2, 3]
+    l[flint.fmpz(1)] = -2
+    assert l == [1, -2, 3]
+    d = {flint.fmpz(2): 3}
+    d[flint.fmpz(2)] = -1
+
+    assert d == {flint.fmpz(2): -1}
     ctx.pretty = False
     assert repr(flint.fmpz(0)) == "fmpz(0)"
     assert repr(flint.fmpz(-27)) == "fmpz(-27)"
     ctx.pretty = True
     assert repr(flint.fmpz(0)) == "0"
     assert repr(flint.fmpz(-27)) == "-27"
-    assert bool(flint.fmpz(0)) == False
-    assert bool(flint.fmpz(1)) == True
     bigstr = '1' * 100
     big = flint.fmpz(bigstr)
     assert big.str() == bigstr
@@ -635,6 +717,7 @@ def test_fmpq():
     assert 0 == Q(0)
     assert Q(2) != 1
     assert 1 != Q(2)
+    assert Q(1) != ()
     assert Q(1,2) != 1
     assert Q(2,3) == Q(flint.fmpz(2),long(3))
     assert Q(-2,-4) == Q(1,2)
@@ -672,6 +755,10 @@ def test_fmpq():
     # XXX: This should NotImplementedError or something.
     assert raises(lambda: pow(Q(1,2),2,3), AssertionError)
 
+    megaz = flint.fmpz(2) ** 8000000
+    megaq = Q(megaz)
+    assert raises(lambda: megaq ** megaz, OverflowError)
+
     assert raises(lambda: Q(1,2) + [], TypeError)
     assert raises(lambda: Q(1,2) - [], TypeError)
     assert raises(lambda: Q(1,2) * [], TypeError)
@@ -700,14 +787,16 @@ def test_fmpq():
     assert (Q(1,2) >= Q(1,2)) is True
     assert raises(lambda: Q(1,2) > [], TypeError)
     assert raises(lambda: [] < Q(1,2), TypeError)
+
     ctx.pretty = False
     assert repr(Q(-2,3)) == "fmpq(-2,3)"
     assert repr(Q(3)) == "fmpq(3)"
     ctx.pretty = True
     assert str(Q(-2,3)) == "-2/3"
     assert str(Q(3)) == "3"
-    assert Q(2,3).p == Q(2,3).numer() == 2
-    assert Q(2,3).q == Q(2,3).denom() == 3
+
+    assert Q(2,3).p == Q(2,3).numer() == Q(2,3).numerator == 2
+    assert Q(2,3).q == Q(2,3).denom() == Q(2,3).denominator == 3
     assert +Q(5,7) == Q(5,7)
     assert -Q(5,7) == Q(-5,7)
     assert -Q(-5,7) == Q(5,7)
@@ -721,12 +810,26 @@ def test_fmpq():
     assert Q(-5,3).floor() == flint.fmpz(-2)
     assert Q(5,3).ceil() == flint.fmpz(2)
     assert Q(-5,3).ceil() == flint.fmpz(-1)
-    # XXX: Need __floor__ etc.
-    #
-    # assert math.floor(Q(5,3)) == flint.fmpz(1)
-    # assert math.ceil(Q(5,3)) == flint.fmpz(2)
-    # assert math.trunc(Q(5,3)) == flint.fmpz(2)
-    # assert round(Q(5,3)) == 2
+
+    assert int(Q(5,3)) == flint.fmpz(1)
+    assert math.floor(Q(5,3)) == flint.fmpz(1)
+    assert math.ceil(Q(5,3)) == flint.fmpz(2)
+    assert math.trunc(Q(5,3)) == flint.fmpz(1)
+    assert round(Q(5,3)) == 2
+
+    assert int(Q(-5,3)) == flint.fmpz(-1)
+    assert math.floor(Q(-5,3)) == flint.fmpz(-2)
+    assert math.ceil(Q(-5,3)) == flint.fmpz(-1)
+    assert math.trunc(Q(-5,3)) == flint.fmpz(-1)
+    assert round(Q(-5,3)) == -2
+
+    assert type(round(Q(5,3))) is flint.fmpz
+    assert type(round(Q(5,3), 0)) is flint.fmpq
+    assert type(round(Q(5,3), 1)) is flint.fmpq
+    assert round(Q(100,3), 2) == Q(3333,100)
+    assert round(Q(100,3), 0) == Q(33,1)
+    assert round(Q(100,3), -1) == Q(30,1)
+    assert round(Q(100,3), -2) == Q(0)
 
     d = {}
     d[Q(1,2)] = 3
@@ -1429,6 +1532,17 @@ def test_arb():
     assert A(3) != A(2)
     assert not (A("1.1") == A("1.1"))
 
+def test_pickling():
+    objects = [
+        flint.fmpz(1),
+        flint.fmpq(1,2),
+        # XXX: Add pickling for everything else
+    ]
+    for obj in objects:
+        s = pickle.dumps(obj)
+        obj2 = pickle.loads(s)
+        assert obj == obj2
+
 if __name__ == "__main__":
     sys.stdout.write("test_pyflint..."); test_pyflint(); print("OK")
     sys.stdout.write("test_fmpz..."); test_fmpz(); print("OK")