Create PY2 constant to simplify compatibility decisions (sybrenstuvel#82)

adamantike · sybrenstuvel · commit 81f0e95dd17b · 2017-01-15T16:30:37.000+01:00
It's about time to get this merged, thanks again!
diff --git 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
@@ -83,11 +83,7 @@ def keygen():
             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 @@ def write_outfile(self, outdata, outname):
                 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
@@ -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
@@ -16,7 +16,6 @@
 
 import unittest
 import struct
-import sys
 
 from rsa._compat import byte, is_bytes, range
 
