Create PY2 constant to simplify compatibility decisions (#82)
It's about time to get this merged, thanks again!
diff --git a/rsa/_compat.py b/rsa/_compat.py
index 6824bdb..1e51368 100644
--- a/rsa/_compat.py
+++ b/rsa/_compat.py
@@ -26,6 +26,8 @@
MAX_INT32 = (1 << 31) - 1
MAX_INT16 = (1 << 15) - 1
+PY2 = sys.version_info[0] == 2
+
# Determine the word size of the processor.
if MAX_INT == MAX_INT64:
# 64-bit processor.
@@ -37,19 +39,24 @@
# Else we just assume 64-bit processor keeping up with modern times.
MACHINE_WORD_SIZE = 64
-# Range generator.
-try:
- # < Python3
+if PY2:
+ integer_types = (int, long)
range = xrange
-except NameError:
- # Python3
+else:
+ integer_types = (int, )
range = range
-# ``long`` is no more. Do type detection using this instead.
-try:
- integer_types = (int, long)
-except NameError:
- integer_types = (int,)
+
+def write_to_stdout(data):
+ """Writes bytes to stdout
+
+ :type data: bytes
+ """
+ if PY2:
+ sys.stdout.write(data)
+ else:
+ # On Py3 we must use the buffer interface to write bytes.
+ sys.stdout.buffer.write(data)
def is_bytes(obj):
diff --git a/rsa/cli.py b/rsa/cli.py
index 7419780..6450af4 100644
--- a/rsa/cli.py
+++ b/rsa/cli.py
@@ -83,11 +83,7 @@
outfile.write(data)
else:
print('Writing private key to stdout', file=sys.stderr)
- if sys.version_info[0] >= 3:
- # on Py3 we must use the buffer interface to write bytes.
- sys.stdout.buffer.write(data)
- else:
- sys.stdout.write(data)
+ rsa._compat.write_to_stdout(data)
class CryptoOperation(object):
@@ -193,11 +189,7 @@
outfile.write(outdata)
else:
print('Writing output to stdout', file=sys.stderr)
- if sys.version_info[0] >= 3:
- # on Py3 we must use the buffer interface to write bytes.
- sys.stdout.buffer.write(outdata)
- else:
- sys.stdout.write(outdata)
+ rsa._compat.write_to_stdout(outdata)
class EncryptOperation(CryptoOperation):
diff --git a/tests/test_cli.py b/tests/test_cli.py
index 4ae8da3..7ce57eb 100644
--- a/tests/test_cli.py
+++ b/tests/test_cli.py
@@ -15,25 +15,23 @@
import rsa
import rsa.cli
import rsa.util
+from rsa._compat import PY2
-if sys.version_info[0] < 3:
- def make_buffer():
+
+def make_buffer():
+ if PY2:
return BytesIO()
+ buf = StringIO()
+ buf.buffer = BytesIO()
+ return buf
- def get_bytes_out(out):
- # Python 2.x writes 'str' to stdout:
+def get_bytes_out(out):
+ if PY2:
+ # Python 2.x writes 'str' to stdout
return out.getvalue()
-else:
- def make_buffer():
- buf = StringIO()
- buf.buffer = BytesIO()
- return buf
-
-
- def get_bytes_out(out):
- # Python 3.x writes 'bytes' to stdout.buffer:
- return out.buffer.getvalue()
+ # Python 3.x writes 'bytes' to stdout.buffer
+ return out.buffer.getvalue()
@contextmanager
diff --git a/tests/test_compat.py b/tests/test_compat.py
index 0013155..a47f890 100644
--- a/tests/test_compat.py
+++ b/tests/test_compat.py
@@ -16,7 +16,6 @@
import unittest
import struct
-import sys
from rsa._compat import byte, is_bytes, range