py: int: tests for **, <<, >> and fixes

ncw · ncw · commit db8950cfaf91 · 2015-06-03T12:41:44.000+01:00
diff --git a/py/bigint.go b/py/bigint.go
@@ -276,6 +276,26 @@ func (a *BigInt) M__rdivmod__(other Object) (Object, Object, error) {
 	return NotImplemented, NotImplemented, nil
 }
 
+// Raise to the power a**b or if m != nil, a**b mod m
+func (a *BigInt) pow(b, m *BigInt) (Object, error) {
+	// -ve power => make float
+	if (*big.Int)(b).Sign() < 0 {
+		if m != nil {
+			return nil, ExceptionNewf(TypeError, "pow() 2nd argument cannot be negative when 3rd argument specified")
+		}
+		fa, err := a.Float()
+		if err != nil {
+			return nil, err
+		}
+		fb, err := b.Float()
+		if err != nil {
+			return nil, err
+		}
+		return fa.M__pow__(fb, None)
+	}
+	return (*BigInt)(new(big.Int).Exp((*big.Int)(a), (*big.Int)(b), (*big.Int)(m))).MaybeInt(), nil
+}
+
 func (a *BigInt) M__pow__(other, modulus Object) (Object, error) {
 	var m *BigInt
 	if modulus != None {
@@ -285,14 +305,14 @@ func (a *BigInt) M__pow__(other, modulus Object) (Object, error) {
 		}
 	}
 	if b, ok := convertToBigInt(other); ok {
-		return (*BigInt)(new(big.Int).Exp((*big.Int)(a), (*big.Int)(b), (*big.Int)(m))).MaybeInt(), nil
+		return a.pow(b, m)
 	}
 	return NotImplemented, nil
 }
 
 func (a *BigInt) M__rpow__(other Object) (Object, error) {
 	if b, ok := convertToBigInt(other); ok {
-		return (*BigInt)(new(big.Int).Exp((*big.Int)(b), (*big.Int)(a), nil)).MaybeInt(), nil
+		return b.pow(a, nil)
 	}
 	return NotImplemented, nil
 }
diff --git a/py/tests/int.py b/py/tests/int.py
@@ -454,14 +454,104 @@ def approxEqual(a, b):
 a ^= -236051837221964566609479775048
 assert a == 737915270162409905551579208985
 
+doc="<< and >>"
+assert (-1 << 10) == -1024
+assert (-1 >> 10) == -1
+assert (3 << 10) == 3072
+assert (3 >> 10) == 0
+assert (-1 << 100) == -1267650600228229401496703205376
+assert (-1267650600228229401496703205376 >> 50) == -1125899906842624
+assert (-1267650600228229401496703205376 >> 90) == -1024
+assert (1 << 100) == 1267650600228229401496703205376
+assert (1267650600228229401496703205376 >> 50) == 1125899906842624
+assert (1267650600228229401496703205376 >> 90) == 1024
+try:
+    1 << 100000000000000000000000000000000000
+except OverflowError:
+    pass
+else:
+    assert False, "Overflow error not raised"
+
+doc="<<= and >>="
+a = -1
+a <<= 10
+assert a == -1024
+a = -1
+a >>= 10
+assert a == -1
+a = 3
+a <<= 10
+assert a == 3072
+a = 3
+a >>= 10
+assert a == 0
+a = -1
+a <<= 100
+assert a == -1267650600228229401496703205376
+a = -1267650600228229401496703205376
+a >>= 50
+assert a == -1125899906842624
+a = -1267650600228229401496703205376
+a >>= 90
+assert a == -1024
+a = 1
+a <<= 100
+assert a == 1267650600228229401496703205376
+a = 1267650600228229401496703205376
+a >>= 50
+assert a == 1125899906842624
+a = 1267650600228229401496703205376
+a >>= 90
+assert a == 1024
+try:
+    a = 1
+    a <<= 100000000000000000000000000000000000
+except OverflowError:
+    pass
+else:
+    assert False, "Overflow error not raised"
+
+doc="**"
+assert (2**10) == 1024
+assert ((-2)**10) == 1024
+approxEqual(2**(-10), 0.0009765625)
+approxEqual((-2)**(-10), 0.0009765625)
+assert (3**100) == 515377520732011331036461129765621272702107522001
+approxEqual(515377520732011331036461129765621272702107522001**(-1), 1.9403252174826328e-48)
+assert ((-1)**100000000000000000000) == 1
+assert ((-1)**100000000000000000001) == -1
+
+doc="**="
+a = 2
+a **= 10
+assert a == 1024
+a = -2
+a **= 10
+assert a == 1024
+a = 2
+a **= -10
+approxEqual(a, 0.0009765625)
+a = -2
+a **= -10
+approxEqual(a, 0.0009765625)
+
+a = 3
+a **= 100
+assert a == 515377520732011331036461129765621272702107522001
+a = 515377520732011331036461129765621272702107522001
+a **= -1
+approxEqual(a, 1.9403252174826328e-48)
+a = -1
+a **= 100000000000000000000
+assert a == 1
+a = -1
+a **= 100000000000000000001
+assert a == -1
+
 # FIXME
-# powmod
+# pow
 # divmod
-# <<
-# >>
 # round
-# **
-# **=
 
 doc="finished"
 
