From fd677e8743953ecd7c009f25943366696568f4e6 Mon Sep 17 00:00:00 2001
From: Daylin Morgan <daylinmorgan@gmail.com>
Date: Thu, 11 May 2023 07:29:55 -0500
Subject: [PATCH] refactor: use themedb

---
 scripts/theme-showcase-gen |  2 +-
 src/yartsu/cli.py          | 15 +++++------
 src/yartsu/themes.py       | 53 +++++++++++++++++++++++++-------------
 3 files changed, 43 insertions(+), 27 deletions(-)

diff --git a/scripts/theme-showcase-gen b/scripts/theme-showcase-gen
index ae2faad..6562cb5 100755
--- a/scripts/theme-showcase-gen
+++ b/scripts/theme-showcase-gen
@@ -2,7 +2,7 @@
 
 THEME_MD="./docs/themes.md"
 
-themes=$(yartsu --list-themes | tail -n +2 | awk '{print $2}')
+themes=$(yartsu --list-themes | tail -n +5 | cut -d' ' -f 3)
 
 newline() {
 	echo >>"$THEME_MD"
diff --git a/src/yartsu/cli.py b/src/yartsu/cli.py
index bfa2ffb..7c605ab 100644
--- a/src/yartsu/cli.py
+++ b/src/yartsu/cli.py
@@ -1,4 +1,3 @@
-import os
 import sys
 import textwrap
 from argparse import SUPPRESS, FileType
@@ -13,9 +12,9 @@ from ._version import __version__
 from .argparse import ArgumentParser
 from .console import Console
 from .term import term
-from .themes import THEMES
+from .themes import ThemeDB
 
-DEFAULT_THEME = os.getenv("YARTSU_THEME", "cat-mocha")
+themes = ThemeDB()
 
 
 def get_parser() -> ArgumentParser:
@@ -56,7 +55,7 @@ def get_parser() -> ArgumentParser:
         "--theme",
         help="theme to use for highlighting [default: %(default)s]",
         type=str,
-        default=DEFAULT_THEME,
+        default=themes.default,
     )
     parser.add_argument(
         "--list-themes", help="list available themes", action="store_true"
@@ -71,8 +70,7 @@ def main() -> None:
     console = Console(record=True)
 
     if args.list_themes:
-        term.print("Available themes:")
-        term.print("\n".join(["  - " + theme for theme in THEMES]))
+        themes.list()
         sys.exit(0)
 
     if args.cmd and args.input or not (args.cmd or args.input or args.demo):
@@ -85,7 +83,8 @@ def main() -> None:
         parser.print_help()
         sys.exit(1)
 
-    if args.theme not in THEMES:
+    # TODO: move this error somewhere else
+    if args.theme not in themes.themes:
         term.print(f"[ThemeError]: {args.theme} is not a valid theme", err=True)
         sys.exit(1)
 
@@ -126,7 +125,7 @@ def main() -> None:
     console.save_svg(
         args.output,
         title=title,
-        theme=THEMES[args.theme],
+        theme=themes.themes[args.theme],
         code_format=CONSOLE_SVG_FORMAT,
     )
 
diff --git a/src/yartsu/themes.py b/src/yartsu/themes.py
index db0a9eb..4546f96 100644
--- a/src/yartsu/themes.py
+++ b/src/yartsu/themes.py
@@ -1,8 +1,11 @@
 import json
+import os
 import sys
 from importlib.resources import files
 
+from rich import box
 from rich.color import parse_rgb_hex
+from rich.table import Table
 from rich.terminal_theme import (
     DIMMED_MONOKAI,
     MONOKAI,
@@ -61,23 +64,37 @@ class YartsuTheme(TerminalTheme):
         return cls(background, foreground, colors, bright_colors, src=src)
 
 
-def get_builtin_themes():
-    return (
-        resource.name.split(".")[0]
-        for resource in (files("yartsu") / "themes").iterdir()
-        if resource.is_file()
-    )
+class ThemeDB:
+    def __init__(self):
+        self.default = os.getenv("YARTSU_THEME", "cat-mocha")
+        self.selected = self.default
+        self.themes = {
+            **self._load_yartsu_themes(),
+            **{
+                "dimmed_monokai": DIMMED_MONOKAI,
+                "monokai": MONOKAI,
+                "night-owlish": NIGHT_OWLISH,
+                "rich-default": SVG_EXPORT_THEME,
+            },
+        }
 
+    def _load_yartsu_themes(self):
+        return {
+            name: YartsuTheme.load_theme(name, src="yartsu")
+            for name in sorted(
+                resource.name.split(".")[0]
+                for resource in (files("yartsu") / "themes").iterdir()
+                if resource.is_file()
+            )
+        }
 
-THEMES = {
-    **{
-        name: YartsuTheme.load_theme(name, src="yartsu")
-        for name in get_builtin_themes()
-    },
-    **{
-        "monokai": MONOKAI,
-        "dimmed_monokai": DIMMED_MONOKAI,
-        "night-owlish": NIGHT_OWLISH,
-        "rich-default": SVG_EXPORT_THEME,
-    },
-}
+    def list(self):
+        table = Table(title="Available Themes", box=box.MINIMAL)
+        table.add_column("name")
+        table.add_column("source")
+
+        for name, theme in self.themes.items():
+            source = theme.src if isinstance(theme, YartsuTheme) else "rich"
+            table.add_row(name, source)
+
+        term.print(table)