From 50379e9a2a35446735c323c612587de3932a2aac Mon Sep 17 00:00:00 2001
From: evilchili <evilchili@gmail.com>
Date: Sat, 30 Jul 2022 20:44:16 -0700
Subject: [PATCH] adding combinable tables

---
 pyproject.toml       |  1 +
 rolltable/tables.py  | 49 ++++++++++++++++++++++++++++++++++++++++++--
 tests/test_tables.py | 37 +++++++++++++++++++++++++++++++++
 3 files changed, 85 insertions(+), 2 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index e0612bd..10599b4 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -13,6 +13,7 @@ python = "^3.7"
 typer = "latest"
 rich = "latest"
 pyyaml = "latest"
+pytest = "latest"
 
 
 [tool.poetry.scripts]
diff --git a/rolltable/tables.py b/rolltable/tables.py
index 14f2a89..b940bd6 100644
--- a/rolltable/tables.py
+++ b/rolltable/tables.py
@@ -91,7 +91,7 @@ class RollTable:
         if not self._rows:
             rows = []
             if self.headers:
-                rows.append(self.headers)
+                rows.append(['Roll'] + self.headers)
             if self._collapsed:
                 for line in self._collapsed_rows():
                     rows.append(line)
@@ -172,10 +172,55 @@ class RollTable:
         Return the rows as a single string.
         """
         rows = list(self.rows)
-        str_format = '\t'.join(['{:s}'] * len(rows[0]))
+        str_format = '\t'.join(['{:10s}'] * len(rows[0]))
         return "\n".join([str_format.format(*row) for row in rows])
 
 
+class CombinedTable(RollTable):
+    """
+    Create a table that is a union of other tables.
+    """
+
+    def __init__(self, tables: List[str], die: Optional[int] = 20):
+        self._die = die
+        self._tables = tables
+        self._rows = None
+        self._headers = None
+
+        # reset any cached values
+        for t in self._tables:
+            t._rows = None
+            t._values = None
+            t._collapsed = False
+            t._die = self._die
+
+    @property
+    def tables(self) -> List:
+        return self._tables
+
+    @property
+    def rows(self) -> List:
+        """
+        Compute the rows of the table by concatenating the rows of the individual tables.
+        """
+        if not self._rows:
+
+            # if one table has headers, they must all have them, so fill with empty strings.
+            if sum([1 for t in self.tables if t.headers]) < len(self.tables):
+                for t in self.tables:
+                    if not t.headers:
+                        t._headers = ['.'] * len(t.values[0])
+
+            self._rows = []
+            for i in range(self._die):
+                row = [self.tables[0].rows[i][0]]
+                for x in range(len(self.tables)):
+                    for col in self.tables[x].rows[i][1:]:
+                        row.append(col)
+                self._rows.append(row)
+        return self._rows
+
+
 if __name__ == '__main__':
     import sys
     print(RollTable(path=sys.argv[1], die=int(sys.argv[2])))
diff --git a/tests/test_tables.py b/tests/test_tables.py
index 051eb12..372e750 100644
--- a/tests/test_tables.py
+++ b/tests/test_tables.py
@@ -50,6 +50,43 @@ option 1:
     -  choice 1
 """
 
+fixture_combined_A = """
+A1:
+  - A choice 1
+  - A choice 2
+  - A choice 3
+A2:
+  - A choice 4
+  - A choice 5
+  - A choice 6
+A3:
+  - A choice 7
+  - A choice 8
+  - A choice 9
+"""
+
+fixture_combined_B = """
+metadata:
+    headers:
+        - HeaderB
+        - HeaderB_Choice
+B1:
+  - B choice 1
+B2:
+  - B choice 2
+B3:
+  - B choice 3
+"""
+
+
+def test_combined_tables():
+    tA = tables.RollTable(fixture_combined_A)
+    tB = tables.RollTable(fixture_combined_B)
+
+    combined = tables.CombinedTable(tables=[tA, tB], die=6)
+    assert 'A1' in str(combined)
+    assert 'B1' in str(combined)
+
 
 def test_table_end_to_end():
     assert str(tables.RollTable(fixture_source))