pyocf: improve random string generator

Set of random characters may be exteded with a custom list.

Signed-off-by: Michal Mielewczyk <michal.mielewczyk@intel.com>
This commit is contained in:
Michal Mielewczyk 2021-03-02 13:52:40 +01:00
parent a81be31dd4
commit 7f862c3080

View File

@ -66,14 +66,14 @@ class RandomGenerator:
class RandomStringGenerator: class RandomStringGenerator:
def __init__(self, len_range=Range(0, 20), count=700): def __init__(self, len_range=Range(0, 20), count=700, extra_chars=[]):
with open("config/random.cfg") as f: with open("config/random.cfg") as f:
self.random = random.Random(int(f.read())) self.random = random.Random(int(f.read()))
self.generator = self.__string_generator(len_range) self.generator = self.__string_generator(len_range, extra_chars)
self.count = count self.count = count
self.n = 0 self.n = 0
def __string_generator(self, len_range): def __string_generator(self, len_range, extra_chars):
while True: while True:
for t in [string.digits, for t in [string.digits,
string.ascii_letters + string.digits, string.ascii_letters + string.digits,
@ -81,7 +81,8 @@ class RandomStringGenerator:
string.ascii_uppercase, string.ascii_uppercase,
string.printable, string.printable,
string.punctuation, string.punctuation,
string.hexdigits]: string.hexdigits,
*extra_chars]:
yield ''.join(random.choice(t) for _ in range( yield ''.join(random.choice(t) for _ in range(
self.random.randint(len_range.min, len_range.max) self.random.randint(len_range.min, len_range.max)
)) ))