From 81f0e95dd17bbe99df2ac4958bd9511dd05c788f Mon Sep 17 00:00:00 2001 From: Michael Manganiello Date: Sun, 15 Jan 2017 12:30:37 -0300 Subject: Create PY2 constant to simplify compatibility decisions (#82) It's about time to get this merged, thanks again! --- tests/test_cli.py | 26 ++++++++++++-------------- tests/test_compat.py | 1 - 2 files changed, 12 insertions(+), 15 deletions(-) (limited to 'tests') 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 -- cgit v1.2.1