From 75f6991b8cd3b74482349030d5278d7c4c14277c Mon Sep 17 00:00:00 2001
From: evilchili <evilchili@gmail.com>
Date: Fri, 22 Dec 2023 14:44:50 -0800
Subject: [PATCH] adding support for flattening entries in a datasource

---
 rolltable/tables.py  | 56 +++++++++++++++++++++++++-------------------
 tests/test_tables.py | 13 ++++++----
 2 files changed, 40 insertions(+), 29 deletions(-)

diff --git a/rolltable/tables.py b/rolltable/tables.py
index e085281..8ddc4b9 100644
--- a/rolltable/tables.py
+++ b/rolltable/tables.py
@@ -83,13 +83,13 @@ class DataSource:
         Return a list of random values from the data set, as a list of lists.
         """
         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
-        list of the option, the select item, and the item's value (if any).
+        For a random item or each item in the specified option in the data source,
+        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.
@@ -103,34 +103,42 @@ class DataSource:
             #
             #  >>> 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()))
-            choice = [k] + v
+            if rand:
+                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:
             # 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
-        # of the option name, the key, and the value, flattening the # value if it is also a list.
-        if hasattr(choice, 'keys'):
-            for (k, v) in choice.items():
-                if type(v) is list:
-                    flattened.extend([k, *v])
-                else:
-                    flattened.extend([k, v])
-            return flattened
+        for choice in choices:
+            # If the randomly-selected choice is a dict, choose a random item and return a list consisting
+            # of the option name, the key, and the value, flattening the # value if it is also a list.
+            if hasattr(choice, 'keys'):
+                for (k, v) in choice.items():
+                    if type(v) is list:
+                        flattened.extend([k, *v])
+                    else:
+                        flattened.extend([k, v])
+                continue
 
-        # if the member is a list, return the flattened list
-        if type(choice) is list:
-            flattened.extend(choice)
-            return flattened
+            # if the member is a list, return the flattened list
+            if type(choice) is list:
+                flattened.extend(choice)
+                continue
 
-        # otherwise, return a list consisting of option and choice
-        flattened.append(choice)
+            # otherwise, return a list consisting of option and choice
+            flattened.append(choice)
         return flattened
 
-
 class RollTable:
     """
     Generate a roll table using weighted distributions of random options.
diff --git a/tests/test_tables.py b/tests/test_tables.py
index bd7a6a2..b094014 100644
--- a/tests/test_tables.py
+++ b/tests/test_tables.py
@@ -187,11 +187,14 @@ def test_no_options():
     assert str(t)
 
 
-def test_yaml():
-    assert tables.RollTable([fixture_no_options]).as_yaml()
-    assert tables.RollTable([fixture_one_choice]).as_yaml()
-    assert tables.RollTable([fixture_metadata + fixture_source]).as_yaml()
-    assert tables.RollTable([fixture_source]).as_yaml()
+@pytest.mark.parametrize('table', [
+    tables.RollTable([fixture_no_options]),
+    tables.RollTable([fixture_one_choice]),
+    tables.RollTable([fixture_metadata + fixture_source]),
+    tables.RollTable([fixture_source]),
+])
+def test_yaml(table):
+    assert table.as_yaml()
 
 
 def test_text():