adding support for flattening entries in a datasource

This commit is contained in:
evilchili 2023-12-22 14:44:50 -08:00
parent fdbd96e9b1
commit 75f6991b8c
2 changed files with 40 additions and 29 deletions

View File

@ -83,13 +83,13 @@ class DataSource:
Return a list of random values from the data set, as a list of lists. Return a list of random values from the data set, as a list of lists.
""" """
return [ return [
self._random_choice_from_option(option) for option in self.random_frequencies(count) self.get_entries(option, rand=True) for option in self.random_frequencies(count)
] ]
def _random_choice_from_option(self, option: str) -> list: def get_entries(self, option, rand: bool = True) -> list:
""" """
Select a random item from the specified option in the data source, and return a flattened For a random item or each item in the specified option in the data source,
list of the option, the select item, and the item's value (if any). return a flattened list of the option, the select item, and the item's value (if any).
""" """
# If there is no data for the specified option, stop now. # If there is no data for the specified option, stop now.
@ -103,34 +103,42 @@ class DataSource:
# #
# >>> self.data[option] == {'One': ['bar', 'baz'], 'Two': ['qaz', 'qux']} # >>> self.data[option] == {'One': ['bar', 'baz'], 'Two': ['qaz', 'qux']}
# #
# choice might then be: ['One', 'bar', 'baz'] # choices might then be: ['One', 'bar', 'baz']
# #
k, v = random.choice(list(self.data[option].items())) if rand:
choice = [k] + v k, v = random.choice(list(self.data[option].items()))
choices = [[k] + v]
else:
choices = [
[k] + v for k, v in list(self.data[option].items())
]
else: else:
# If the option is either a list or a string, just select it. # If the option is either a list or a string, just select it.
choice = random.choice(self.data[option]) if rand:
choices = [random.choice(self.data[option])]
else:
choices = self.data[option]
# If the randomly-selected choice is a dict, choose a random item and return a list consisting for choice in choices:
# of the option name, the key, and the value, flattening the # value if it is also a list. # If the randomly-selected choice is a dict, choose a random item and return a list consisting
if hasattr(choice, 'keys'): # of the option name, the key, and the value, flattening the # value if it is also a list.
for (k, v) in choice.items(): if hasattr(choice, 'keys'):
if type(v) is list: for (k, v) in choice.items():
flattened.extend([k, *v]) if type(v) is list:
else: flattened.extend([k, *v])
flattened.extend([k, v]) else:
return flattened flattened.extend([k, v])
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: if type(choice) is list:
flattened.extend(choice) flattened.extend(choice)
return flattened 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 flattened return flattened
class RollTable: class RollTable:
""" """
Generate a roll table using weighted distributions of random options. Generate a roll table using weighted distributions of random options.

View File

@ -187,11 +187,14 @@ def test_no_options():
assert str(t) assert str(t)
def test_yaml(): @pytest.mark.parametrize('table', [
assert tables.RollTable([fixture_no_options]).as_yaml() tables.RollTable([fixture_no_options]),
assert tables.RollTable([fixture_one_choice]).as_yaml() tables.RollTable([fixture_one_choice]),
assert tables.RollTable([fixture_metadata + fixture_source]).as_yaml() tables.RollTable([fixture_metadata + fixture_source]),
assert tables.RollTable([fixture_source]).as_yaml() tables.RollTable([fixture_source]),
])
def test_yaml(table):
assert table.as_yaml()
def test_text(): def test_text():