diff --git a/rolltable/tables.py b/rolltable/tables.py index 1021dd5..1862ac5 100644 --- a/rolltable/tables.py +++ b/rolltable/tables.py @@ -135,17 +135,23 @@ class RollTable: values.append([option]) continue if hasattr(ds.data[option], 'keys'): - rand_key = random.choice(list(ds.data[option].keys())) - choice = [rand_key, *ds.data[option][rand_key]] + k, v = random.choice(list(ds.data[option].items())) + choice = [k] + v else: choice = random.choice(ds.data[option]) if hasattr(choice, 'keys'): c = [option] for (k, v) in choice.items(): - c.extend([k, v]) + if type(v) is list: + c.extend([k, *v]) + else: + c.extend([k, v]) values.append(c) else: - values.append([option, choice]) + if type(choice) is list: + values.append([option, *choice]) + else: + values.append([option, choice]) return sorted(values) ds_values = [values_from_datasource(t) for t in self._data] diff --git a/tests/test_tables.py b/tests/test_tables.py index ba234a7..bd7a6a2 100644 --- a/tests/test_tables.py +++ b/tests/test_tables.py @@ -1,3 +1,5 @@ +import pytest + from rolltable import tables fixture_metadata = """ @@ -104,7 +106,7 @@ metadata: - two - three - four -Category: +list: - foo: - bar - baz @@ -121,7 +123,7 @@ metadata: - two - three - four -Category: +dict: foo: - bar - baz @@ -133,9 +135,10 @@ Category: """] -def test_lists_and_dicts(): - t = tables.RollTable(fixture_lists_and_dicts, die=1) - assert str(t) +@pytest.mark.parametrize('fixture', fixture_lists_and_dicts) +def test_lists_and_dicts(fixture): + t = tables.RollTable([fixture], die=1) + assert(str(t)) def test_combined_tables():