From 76ed17a65039c21498d8cececd5f1443a9b317db Mon Sep 17 00:00:00 2001 From: evilchili Date: Sat, 6 Jan 2024 12:18:31 -0800 Subject: [PATCH] fix datasource.random_values to return the correct number of values --- random_sets/datasources.py | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/random_sets/datasources.py b/random_sets/datasources.py index 0075d2a..4d288a1 100644 --- a/random_sets/datasources.py +++ b/random_sets/datasources.py @@ -122,9 +122,9 @@ class DataSource: """ # If there is no data for the specified option, stop now. - flattened = [option] + flattened = [] if not self.data[option]: - return flattened + return [option] if hasattr(self.data[option], 'keys'): # 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'): for (k, v) in choice.items(): if type(v) is list: - flattened.extend([k, *v]) + flattened.append([option, k, *v]) else: - flattened.extend([k, v]) + flattened.append([option, k, v]) continue # if the member is a list, return the flattened list - if type(choice) is list: + elif type(choice) is list: flattened.extend(choice) continue # otherwise, return a list consisting of option and choice flattened.append(choice) + + # Return all randomized values or just 1. + if rand: + return random.choice(flattened) return flattened