From 75ce2498e8c98ec3fd6583985990a56eeb03f2b8 Mon Sep 17 00:00:00 2001 From: Shadab Zafar Date: Thu, 2 Jun 2016 13:34:18 +0530 Subject: [PATCH] Simplify rand_byte by creating a special case for Python 2 --- pathod/language/generators.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/pathod/language/generators.py b/pathod/language/generators.py index e736e0437..20bb7ae1f 100644 --- a/pathod/language/generators.py +++ b/pathod/language/generators.py @@ -2,6 +2,8 @@ import string import random import mmap +import six + DATATYPES = dict( ascii_letters=string.ascii_letters.encode(), ascii_lowercase=string.ascii_lowercase.encode(), @@ -47,7 +49,9 @@ def rand_byte(chars): """ # bytearray has consistent behaviour on both Python 2 and 3 # while bytes does not - return bytes(bytearray([random.choice(chars)])) + if six.PY2: + return random.choice(chars) + return bytes([random.choice(chars)]) class RandomGenerator(object):