From 590b67639ebf16638854eb05665601b6d2bb48f1 Mon Sep 17 00:00:00 2001
From: Laurent Peuch <cortex@worlddomination.be>
Date: Tue, 22 Jun 2021 03:52:44 +0200
Subject: [PATCH] [enh] integrate rich tables in moulinette

---
 moulinette/__init__.py       | 82 ++++++++++++++++++++++++++++++++++++
 moulinette/interfaces/cli.py |  8 +++-
 2 files changed, 89 insertions(+), 1 deletion(-)

diff --git a/moulinette/__init__.py b/moulinette/__init__.py
index b4f2e06e..b73eb7b9 100755
--- a/moulinette/__init__.py
+++ b/moulinette/__init__.py
@@ -1,7 +1,9 @@
 # -*- coding: utf-8 -*-
 
 from rich import traceback
+from rich.table import Table as RichTable, box
 from rich.console import Console
+from rich.style import Style
 
 from moulinette.core import (
     MoulinetteError,
@@ -61,6 +63,86 @@ def _format_exception():
 
 console.format_exception = _format_exception
 
+
+class Table:
+    def __init__(self, data, columns=None, title=None):
+        self.data = data
+        self.columns = columns
+        self.title = title
+
+    def print(self):
+        if not self.data:
+            return
+
+        assert len(self.data.keys()) == 1, self.data
+
+        table = RichTable(show_header=True)
+        table_name = list(self.data.keys())[0]
+        table.title = f"[bold]{self.title if self.title else table_name}[/]"
+        table.title_style = Style(color="white", bgcolor=None)
+        table.row_styles = ["none", "dim"]
+        table.box = box.SIMPLE_HEAD
+        table.border_style = "bright_yellow"
+
+        self._fill_table(table, table_name)
+
+        console.print(table)
+
+    def _fill_table(self, table, table_name):
+        if len(self.data[table_name]) == 0:
+            return
+
+        if self.columns is None:
+            self.columns = sorted(self.data[table_name][0].keys())
+
+        for column in self.columns:
+            if isinstance(column, dict):
+                pass
+            else:
+                table.add_column(column.replace("_", " ").title())
+
+        for row in self.data[table_name]:
+            values = []
+            for column in self.columns:
+                values.append(row.get(column, ""))
+
+            table.add_row(*values)
+
+
+class TableForDict(Table):
+    key = object()  # this is a flag value to say "key of the dictionary"
+
+    def _fill_table(self, table, table_name):
+        if len(self.data[table_name]) == 0:
+            return
+
+        if self.columns is None:
+            an_item = list(self.data[table_name].keys())[0]
+            self.columns = [self.key] + sorted(self.data[table_name][an_item].keys())
+
+        for column in self.columns:
+            if isinstance(column, dict):
+                header = column.get("header", "")
+                table.add_column(str(header))
+            else:
+                if column == self.key:
+                    table.add_column()
+                else:
+                    table.add_column(column.replace("_", " ").title())
+
+        for key, values in self.data[table_name].items():
+            row_values = [key]
+            for column in self.columns:
+                if isinstance(column, dict):
+                    column = column["key"]
+
+                if column != self.key:
+                    row_values.append(str(values.get(column, "")))
+
+            table.add_row(*row_values)
+
+
+
 # Package functions
 
 
diff --git a/moulinette/interfaces/cli.py b/moulinette/interfaces/cli.py
index f21cb499..55fe9077 100644
--- a/moulinette/interfaces/cli.py
+++ b/moulinette/interfaces/cli.py
@@ -11,7 +11,7 @@ from datetime import date, datetime
 
 import argcomplete
 
-from moulinette import msignals, m18n
+from moulinette import msignals, m18n, Table
 from moulinette.actionsmap import ActionsMap
 from moulinette.core import MoulinetteError, MoulinetteValidationError
 from moulinette.interfaces import (
@@ -512,11 +512,17 @@ class Interface(BaseInterface):
                 import json
                 from moulinette.utils.serialize import JSONExtendedEncoder
 
+                # retro compat situation
+                if isinstance(ret, Table):
+                    ret = ret.data
+
                 print(json.dumps(ret, cls=JSONExtendedEncoder))
             else:
                 plain_print_dict(ret)
         elif isinstance(ret, dict):
             pretty_print_dict(ret)
+        elif isinstance(ret, Table):
+            ret.print()
         else:
             print(ret)