Use long depending on python version

4 files changed, 23 insertions, 12 deletions

diff --git a/openid/dh.py b/openid/dh.py
index ab9f984..28ff403 100644
--- a/openid/dh.py
+++ b/openid/dh.py
@@ -4,6 +4,12 @@ import six
 
 from openid import cryptutil
 
+if six.PY2:
+    long_int = long
+else:
+    assert six.PY3
+    long_int = int
+
 
 def _xor(a_b):
     # Python 2 only
@@ -35,8 +41,8 @@ class DiffieHellman(object):
         return cls(cls.DEFAULT_MOD, cls.DEFAULT_GEN)
 
     def __init__(self, modulus, generator):
-        self.modulus = long(modulus)
-        self.generator = long(generator)
+        self.modulus = long_int(modulus)
+        self.generator = long_int(generator)
 
         self._setPrivate(cryptutil.randrange(1, modulus - 1))
 
diff --git a/openid/server/server.py b/openid/server/server.py
index 701fbdf..2d251d4 100644
--- a/openid/server/server.py
+++ b/openid/server/server.py
@@ -318,7 +318,7 @@ class DiffieHellmanSHA1ServerSession(object):
 
     @ivar consumer_pubkey: The public key sent by the consumer in the
         associate request
-    @type consumer_pubkey: long
+    @type consumer_pubkey: int, long in Python 2
 
     @see: U{OpenID Specs, Mode: associate
         <http://openid.net/specs.bml#mode-associate>}
diff --git a/openid/test/test_cryptutil.py b/openid/test/test_cryptutil.py
index 5a86611..248c100 100644
--- a/openid/test/test_cryptutil.py
+++ b/openid/test/test_cryptutil.py
@@ -12,6 +12,11 @@ from openid import cryptutil
 
 # Most of the purpose of this test is to make sure that cryptutil can
 # find a good source of randomness on this machine.
+if six.PY2:
+    long_int = long
+else:
+    assert six.PY3
+    long_int = int
 
 
 class TestRandRange(unittest.TestCase):
@@ -29,13 +34,13 @@ class TestRandRange(unittest.TestCase):
 
         a = cryptutil.randrange(2 ** 128)
         b = cryptutil.randrange(2 ** 128)
-        assert isinstance(a, long)
-        assert isinstance(b, long)
+        assert isinstance(a, long_int)
+        assert isinstance(b, long_int)
         assert b != a
 
         # Make sure that we can generate random numbers that are larger
         # than platform int size
-        cryptutil.randrange(long(sys.maxsize) + 1)
+        cryptutil.randrange(long_int(sys.maxsize) + 1)
 
 
 class TestLongBinary(unittest.TestCase):
@@ -46,7 +51,7 @@ class TestLongBinary(unittest.TestCase):
         for iteration in range(500):
             n = 0
             for i in range(10):
-                n += long(random.randrange(MAX))
+                n += long_int(random.randrange(MAX))
 
             s = cryptutil.longToBinary(n)
             assert isinstance(s, six.binary_type)
@@ -79,7 +84,7 @@ class TestLongToBase64(unittest.TestCase):
         try:
             for line in f:
                 parts = line.strip().split(' ')
-                assert parts[0] == cryptutil.longToBase64(long(parts[1]))
+                assert parts[0] == cryptutil.longToBase64(long_int(parts[1]))
         finally:
             f.close()
 
@@ -92,6 +97,6 @@ class TestBase64ToLong(unittest.TestCase):
         try:
             for line in f:
                 parts = line.strip().split(' ')
-                assert long(parts[1]) == cryptutil.base64ToLong(parts[0])
+                assert long_int(parts[1]) == cryptutil.base64ToLong(parts[0])
         finally:
             f.close()
diff --git a/openid/test/test_dh.py b/openid/test/test_dh.py
index c9a2c56..ddd84f0 100644
--- a/openid/test/test_dh.py
+++ b/openid/test/test_dh.py
@@ -6,7 +6,7 @@ import unittest
 
 import six
 
-from openid.dh import DiffieHellman, strxor
+from openid.dh import DiffieHellman, long_int, strxor
 
 
 class TestStrXor(unittest.TestCase):
@@ -72,8 +72,8 @@ class TestDiffieHellman(unittest.TestCase):
         try:
             for line in f:
                 parts = line.strip().split(' ')
-                dh._setPrivate(long(parts[0]))
+                dh._setPrivate(long_int(parts[0]))
 
-                assert dh.public == long(parts[1])
+                assert dh.public == long_int(parts[1])
         finally:
             f.close()