From 2b3e82114ff19fe8a28cb8d1cdb7b5fb55c8c111 Mon Sep 17 00:00:00 2001
From: evilchili <evilchili@gmail.com>
Date: Sat, 6 Jan 2024 12:21:30 -0800
Subject: [PATCH] fixing table output and moving it to types.RollTable

---
 rolltable/cli.py   |  8 ++------
 rolltable/types.py | 30 +++++++++++++++++++-----------
 2 files changed, 21 insertions(+), 17 deletions(-)

diff --git a/rolltable/cli.py b/rolltable/cli.py
index 0694bd8..36e50b7 100644
--- a/rolltable/cli.py
+++ b/rolltable/cli.py
@@ -51,13 +51,9 @@ def create(
     if output == OUTPUT_FORMATS.yaml:
         print(rt.as_yaml())
     elif output == OUTPUT_FORMATS.markdown:
-        print(rt.as_markdown)
+        print(rt.as_markdown())
     else:
-        rows = rt.rows if collapsed else rt.expanded_rows
-        table = Table(*rows[0], width=width)
-        for row in rows[1:]:
-            table.add_row(*row)
-        print(table)
+        print(rt.as_table(width=width, expanded=not collapsed))
 
 
 if __name__ == '__main__':
diff --git a/rolltable/types.py b/rolltable/types.py
index a228367..ca0a8c5 100644
--- a/rolltable/types.py
+++ b/rolltable/types.py
@@ -4,6 +4,8 @@ from collections.abc import Iterable
 from typing import Optional, List, Union
 from random_sets.datasources import DataSource
 
+import rich.table
+
 
 class RollTable:
     """
@@ -40,16 +42,6 @@ class RollTable:
         self._generated_values = None
         self._config()
 
-    def as_yaml(self, expanded=False) -> dict:
-        struct = {}
-        for row in self.rows[1:]:
-            struct[row[0]] = {}
-            # 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, sort_keys=False)
-
     @property
     def datasources(self) -> List:
         return self._data
@@ -106,10 +98,26 @@ class RollTable:
             self._rows.append(self._column_filter([f'd{face+1}'] + row))
         return self._rows
 
-    @property
     def as_markdown(self) -> str:
         return Table(self.rows).markdown()
 
+    def as_yaml(self, expanded: bool = False) -> dict:
+        struct = {}
+        for row in self.rows[1:]:
+            struct[row[0]] = {}
+            # 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, sort_keys=False)
+
+    def as_table(self, width: int = 120, expanded: bool = False) -> str:
+        rows = self.expanded_rows if expanded else self.rows
+        table = rich.table.Table(*rows[0], width=width)
+        for row in rows[1:]:
+            table.add_row(*row)
+        return table
+
     def _config(self):
         """
         Parse data sources, generate headers, and create the column filters