diff --git a/rolltable/tables.py b/rolltable/tables.py index 00d4e64..1021dd5 100644 --- a/rolltable/tables.py +++ b/rolltable/tables.py @@ -134,7 +134,11 @@ class RollTable: if not ds.data[option]: values.append([option]) continue - choice = random.choice(ds.data[option]) + if hasattr(ds.data[option], 'keys'): + rand_key = random.choice(list(ds.data[option].keys())) + choice = [rand_key, *ds.data[option][rand_key]] + else: + choice = random.choice(ds.data[option]) if hasattr(choice, 'keys'): c = [option] for (k, v) in choice.items(): diff --git a/tests/test_tables.py b/tests/test_tables.py index dc1b9b1..ba234a7 100644 --- a/tests/test_tables.py +++ b/tests/test_tables.py @@ -92,26 +92,49 @@ B2: B3: """ -fixture_lists = """ +fixture_lists_and_dicts = [""" # -# one two three four -# foo bar baz quz +# category one two three four +# Category foo bar baz quz # metadata: headers: + - category - one - two - three - four -foo: - - bar: +Category: + - foo: + - bar + - baz + - quz +""", """ +# +# category one two three four +# Category foo bar baz quz +# +metadata: + headers: + - category + - one + - two + - three + - four +Category: + foo: + - bar - baz - quz -""" + bar: + - a + - b + - c +"""] -def test_lists(): - t = tables.RollTable([fixture_lists], die=1) +def test_lists_and_dicts(): + t = tables.RollTable(fixture_lists_and_dicts, die=1) assert str(t)