fix datasource.random_values to return the correct number of values

This commit is contained in:
evilchili 2024-01-06 12:18:31 -08:00
parent 51d9aa02b4
commit 76ed17a650

View File

@ -122,9 +122,9 @@ class DataSource:
""" """
# If there is no data for the specified option, stop now. # If there is no data for the specified option, stop now.
flattened = [option] flattened = []
if not self.data[option]: if not self.data[option]:
return flattened return [option]
if hasattr(self.data[option], 'keys'): if hasattr(self.data[option], 'keys'):
# if the option is a dict, we assume the values are lists; we select a random item # if the option is a dict, we assume the values are lists; we select a random item
@ -151,16 +151,20 @@ class DataSource:
if hasattr(choice, 'keys'): if hasattr(choice, 'keys'):
for (k, v) in choice.items(): for (k, v) in choice.items():
if type(v) is list: if type(v) is list:
flattened.extend([k, *v]) flattened.append([option, k, *v])
else: else:
flattened.extend([k, v]) flattened.append([option, k, v])
continue continue
# if the member is a list, return the flattened list # if the member is a list, return the flattened list
if type(choice) is list: elif type(choice) is list:
flattened.extend(choice) flattened.extend(choice)
continue continue
# otherwise, return a list consisting of option and choice # otherwise, return a list consisting of option and choice
flattened.append(choice) flattened.append(choice)
# Return all randomized values or just 1.
if rand:
return random.choice(flattened)
return flattened return flattened