From c8117cb66c5684d0783a2444a3f0d9af0aa9a430 Mon Sep 17 00:00:00 2001
From: evilchili <evilchili@gmail.com>
Date: Sat, 13 Aug 2022 12:32:52 -0700
Subject: [PATCH] fix mismatches in header and column counts

---
 pyproject.toml       |  2 +-
 rolltable/tables.py  | 24 +++++++++++++++++++-----
 tests/test_tables.py | 33 +++++++++++++++++++++++++++++++--
 3 files changed, 51 insertions(+), 8 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 55746a8..c2dd8a8 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = 'dnd-rolltable'
-version = '1.1.3'
+version = '1.1.4'
 license = 'The Unlicense'
 authors = ['Greg Boyington <evilchili@gmail.com>']
 description = 'Generate roll tables using weighted random distributions'
diff --git a/rolltable/tables.py b/rolltable/tables.py
index dc21823..fca4ce2 100644
--- a/rolltable/tables.py
+++ b/rolltable/tables.py
@@ -1,5 +1,6 @@
 import yaml
 import random
+from collections.abc import Iterable
 from typing import Optional, List, IO
 
 
@@ -97,8 +98,10 @@ class RollTable:
         struct = {}
         for row in self.rows[1:]:
             struct[row[0]] = {}
-            for idx, col in enumerate(row[1:]):
-                struct[row[0]][self.headers[idx]] = col
+            # pad rows with empty cols as necessary
+            cols = row[1:] + [''] * (len(self.headers) - len(row[1:]))
+            for idx, col in enumerate(cols):
+                struct[row[0]][self.headers[idx] if idx < len(self.headers) else '_'] = col
         return yaml.dump(struct)
 
     @property
@@ -147,8 +150,9 @@ class RollTable:
     @property
     def rows(self) -> List:
         def formatted(lastrow, offset, row, i):
-            fmt = f'd{i}' if offset + 1 == i else f'd{offset+1}-d{i}'
-            return self._column_filter([fmt] + lastrow)
+            thisrow = [f'd{i}' if offset + 1 == i else f'd{offset+1}-d{i}']
+            thisrow += self._flatten(lastrow)
+            return self._column_filter(thisrow)
 
         lastrow = None
         offset = 0
@@ -203,9 +207,19 @@ class RollTable:
                 self._header_excludes.append(i+1)  # +1 to account for the 'Roll' column
 
     def _column_filter(self, row):
-        return [col for (pos, col) in enumerate(row) if pos not in self._header_excludes]
+        cols = [col for (pos, col) in enumerate(row) if pos not in self._header_excludes]
+        # pad the row with empty columns if there are more headers than columns
+        return cols + [''] * (1 + len(self.headers) - len(row))
+
+    def _flatten(self, obj: List) -> List:
+        for member in obj:
+            if isinstance(member, Iterable) and not isinstance(member, (str, bytes)):
+                yield from self._flatten(member)
+            else:
+                yield member
 
     def __repr__(self) -> str:
         rows = list(self.rows)
+        print(rows)
         str_format = '\t'.join(['{:10s}'] * len(rows[0]))
         return "\n".join([str_format.format(*row) for row in rows])
diff --git a/tests/test_tables.py b/tests/test_tables.py
index eb32068..3581423 100644
--- a/tests/test_tables.py
+++ b/tests/test_tables.py
@@ -92,6 +92,28 @@ B2:
 B3:
 """
 
+fixture_lists = """
+#
+# one  two  three  four
+# foo  bar  baz    quz
+#
+metadata:
+  headers:
+    - one
+    - two
+    - three
+    - four
+foo:
+  - bar:
+    - baz
+    - quz
+"""
+
+
+def test_lists():
+    t = tables.RollTable([fixture_lists], die=1)
+    assert str(t)
+
 
 def test_combined_tables():
     combined = tables.RollTable([fixture_combined_A, fixture_combined_B], die=6)
@@ -134,6 +156,13 @@ def test_no_descriptions():
     assert 'option 1' in str(t)
 
 
-def test_yaml():
+def test_no_options():
     t = tables.RollTable([fixture_no_options])
-    print(t.as_yaml)
+    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()