diff options
| -rw-r--r-- | rsa/_compat.py | 27 | ||||
| -rw-r--r-- | rsa/cli.py | 12 | ||||
| -rw-r--r-- | tests/test_cli.py | 26 | ||||
| -rw-r--r-- | tests/test_compat.py | 1 |
4 files changed, 31 insertions, 35 deletions
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_INT64 = (1 << 63) - 1 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: # 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): @@ -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 @@ class CryptoOperation(object): 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 @@ from io import StringIO, BytesIO 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 |