Use rich theme in argparse help menu (#58)

* - Add AstroTheme class in a new theme.py file
- AstroTheme will now hold all color theme related info
- Some private class methods have been privated using the double underscore convention
- A theme was defined for argparse, and implemented by AstroTheme
- Modified conftest.py to add fixture for returning AstroThme() object
- Modified logger() fixture to pass AstroTheme() object to logger.astro_config() method
- update requirements.txt to include rich_argparse package

Author: AustinCullar

requirements.txt

@@ -36,6 +36,7 @@ pytz==2024.2
 regex==2024.9.11
 requests==2.32.3
 rich==13.9.2
+rich-argparse==1.5.2
 rsa==4.9
 six==1.16.0
 tqdm==4.66.5

setup.py

@@ -10,7 +10,8 @@
     py_modules=[
         "src/log",
         "src/astro_db",
-        "src/progress"],
+        "src/progress",
+        "src/theme"],
 
     # Packages required to run the app
     install_requires=[

src/astro.py

@@ -11,6 +11,8 @@
 from data_collection.sentiment import SentimentAnalysis
 from log import AstroLogger
 from astro_db import AstroDB
+from theme import AstroTheme
+from rich_argparse import ArgumentDefaultsRichHelpFormatter
 
 
 def extract_video_id_from_url(url: str) -> str:
@@ -27,26 +29,31 @@ def extract_video_id_from_url(url: str) -> str:
     return video_id
 
 
-def parse_args():
+def parse_args(astro_theme):
     """
     Argument parsing logic. Returns the arguments parsed from the CLI
     """
+    description = "A tool for YouTube data collection."
 
-    parser = argparse.ArgumentParser()
+    parser = argparse.ArgumentParser(description=description,
+                                     formatter_class=ArgumentDefaultsRichHelpFormatter)
 
     parser.add_argument("youtube_url", type=str, help="youtube video URL")
     parser.add_argument("-l", "--log", type=str, choices=['debug', 'info', 'warn', 'error'],
-                        help='Set the logging level')
+                        help='Set the logging level', default='info')
     parser.add_argument("--api-key", type=str, help="YouTube Data API key")
-    parser.add_argument("--db-file", type=str, help="database filename")
+    parser.add_argument("--db-file", type=str, help="database filename", default='astro.db')
     args = parser.parse_args()
 
     return args
 
 
 def main():
+    # load astro color scheme
+    astro_theme = AstroTheme()
+
     # parse arguments
-    args = parse_args()
+    args = parse_args(astro_theme)
     video_id = extract_video_id_from_url(args.youtube_url)
 
     # load environment variables
@@ -60,7 +67,7 @@ def main():
     # set up logging
     logging.setLoggerClass(AstroLogger)
     logger = logging.getLogger(__name__)
-    logger.astro_config(log_level)
+    logger.astro_config(log_level, astro_theme)
 
     logger.info('Collecting video data...')
 

src/log.py

@@ -20,7 +20,7 @@ class AstroLogger(logging.Logger):
     astro_theme: Theme
     progress: AstroProgress
 
-    def astro_config(self, log_level_str: str):
+    def astro_config(self, log_level_str: str, astro_theme):
         """
         Custom logging config.
         """
@@ -29,21 +29,10 @@ def astro_config(self, log_level_str: str):
 
         self.setLevel(self.log_level)
 
-        # define astro color theme
-        self.astro_text_color = 'sky_blue3'
-        self.astro_theme = Theme({
-            "log.message": self.astro_text_color,
-            "logging.level.info": 'blue_violet',
-            "logging.level.warning": "orange_red1",
-            "logging.level.error": "red",
-            "bar.finished": "green",
-            "progress.elapsed": self.astro_text_color,
-            "progress.remaining": self.astro_text_color,
-            "progress.percentage": "green",
-            "table.title": "bold " + self.astro_text_color})
+        self.astro_theme = astro_theme
 
         # create console using the asto theme
-        self.console = Console(theme=self.astro_theme)
+        self.console = self.astro_theme.get_console()
 
         # configure logging
         logging.basicConfig(format='%(message)s',
@@ -52,9 +41,9 @@ def astro_config(self, log_level_str: str):
                                                   console=self.console)])
 
         # suppress google logs
-        self.suppress_logs('google', logging.WARNING)
+        self.__suppress_logs('google', logging.WARNING)
 
-    def suppress_logs(self, name, level):
+    def __suppress_logs(self, name, level):
         """
         Suppress logging of external modules.
         """
@@ -67,7 +56,7 @@ def progress_bar(self, task_str: str, steps: int):
         """
         Display a progress bar on the console.
         """
-        self.progress = AstroProgress(task_str, steps, console=self.console, style=self.astro_text_color)
+        self.progress = AstroProgress(task_str, steps, console=self.console)
         return self.progress
 
     def get_log_level(self, log_level: str) -> int:
@@ -87,14 +76,15 @@ def get_log_level(self, log_level: str) -> int:
 
         return log_level
 
-    def rich_table(self, title=''):
+    def __rich_table(self, title=''):
         """
         Create a rich.Table object using the default project theme/style.
         """
-        table = Table(style=self.astro_text_color,
-                      border_style=self.astro_text_color,
-                      header_style=self.astro_text_color,
-                      row_styles=[self.astro_text_color],
+        text_color = self.astro_theme.get_style()
+        table = Table(style=text_color,
+                      border_style=text_color,
+                      header_style=text_color,
+                      row_styles=[text_color],
                       title=title)
 
         return table
@@ -110,7 +100,7 @@ def print_object(self, obj, title=''):
         if self.log_level > logging.INFO:
             return
 
-        table = self.rich_table(title)
+        table = self.__rich_table(title)
         table.add_column("Attribute")
         table.add_column("Value")
 
@@ -130,7 +120,7 @@ def print_dataframe(self, df, title=''):
         # ensure dataframe contains only string values
         df = df.astype(str)
 
-        table = self.rich_table(title)
+        table = self.__rich_table(title)
 
         # add dataframe columns
         for col in df.columns:

src/progress.py

@@ -15,21 +15,21 @@ class AstroProgress(Progress):
     total_steps: int
     task: int
 
-    def __init__(self, task_str: str, steps: int, console=None, style=None):
+    def __init__(self, task_str: str, steps: int, console=None):
         self.task_str = task_str
         self.total_steps = steps
 
         super().__init__(
             TextColumn("{task.description} [progress.percentage]{task.percentage:>3.0f}%"),
             BarColumn(),
             MofNCompleteColumn(),
-            TextColumn("•", style=style),
+            TextColumn("•", style=console.get_style('log.message')),
             TimeElapsedColumn(),
-            TextColumn("•", style=style),
+            TextColumn("•", style=console.get_style('log.message')),
             TimeRemainingColumn(),
             console=console)
 
-        self.task = super().add_task(f'[{style}]' + task_str, total=self.total_steps)
+        self.task = super().add_task(f"[{console.get_style('log.message')}]" + task_str, total=self.total_steps)
 
     def advance(self, steps: int):
         super().update(self.task, advance=steps)

src/tests/conftest.py

@@ -4,9 +4,10 @@
 import logging
 
 import pandas as pd
+import src.tests.test_api_responses as api_responses
 from unittest.mock import MagicMock
 from src.log import AstroLogger
-import src.tests.test_api_responses as api_responses
+from src.theme import AstroTheme
 
 
 @pytest.fixture(scope='class')
@@ -31,10 +32,15 @@ def api_video_response():
 
 
 @pytest.fixture(scope='class')
-def logger():
+def astro_theme():
+    return AstroTheme()
+
+
+@pytest.fixture(scope='class')
+def logger(astro_theme):
     logging.setLoggerClass(AstroLogger)
     astro_logger = logging.getLogger(__name__)
-    astro_logger.astro_config('debug')
+    astro_logger.astro_config('debug', astro_theme)
     return astro_logger
 
 

src/theme.py

@@ -0,0 +1,49 @@
+"""
+This file defines the project color themes for anything printed to the console.
+"""
+from rich.theme import Theme
+from rich.console import Console
+from rich_argparse import ArgumentDefaultsRichHelpFormatter
+
+
+class AstroTheme:
+    astro_theme: Theme
+    astro_console: Console
+    astro_text_color = 'sky_blue3'
+
+    def __init__(self):
+        self.astro_theme = Theme({
+            "log.message": self.astro_text_color,
+            "logging.level.info": 'blue_violet',
+            "logging.level.warning": "orange_red1",
+            "logging.level.error": "red",
+            "bar.finished": "green",
+            "progress.elapsed": self.astro_text_color,
+            "progress.remaining": self.astro_text_color,
+            "progress.percentage": "green",
+            "table.title": "bold " + self.astro_text_color})
+
+        self.astro_console = Console(theme=self.astro_theme)
+        self.__set_argparse_theme()
+
+    def __set_argparse_theme(self):
+        """
+        Set the theme of the argparse help menu.
+        """
+        ArgumentDefaultsRichHelpFormatter.styles['argparse.args'] = self.astro_text_color  # for positional-arguments
+        ArgumentDefaultsRichHelpFormatter.styles['argparse.groups'] = self.astro_text_color  # for group names
+        ArgumentDefaultsRichHelpFormatter.styles['argparse.help'] = self.astro_text_color  # for argument's help text
+        ArgumentDefaultsRichHelpFormatter.styles['argparse.metavar'] = 'green'  # for metavariables
+        ArgumentDefaultsRichHelpFormatter.styles['argparse.prog'] = 'green'  # for %(prog)s in the usage
+        ArgumentDefaultsRichHelpFormatter.styles['argparse.syntax'] = 'bold'  # for highlights of back-tick quoted text
+        ArgumentDefaultsRichHelpFormatter.styles['argparse.text'] = self.astro_text_color  # for descriptions
+        ArgumentDefaultsRichHelpFormatter.styles['argparse.default'] = 'italic'  # for %(default)s in the help
+
+    def get_style(self):
+        return self.astro_text_color
+
+    def get_theme(self):
+        return self.astro_theme
+
+    def get_console(self):
+        return self.astro_console