Colour

datafaker/interactive/base.py

@@ -9,7 +9,7 @@
 from typing import Any, Optional, Type
 
 import sqlalchemy
-from prettytable import PrettyTable
+from prettytable.colortable import ColorTable, Theme
 from sqlalchemy import Engine, ForeignKey, MetaData, Table, func, or_, select
 from sqlalchemy.exc import DatabaseError, SQLAlchemyError
 from typing_extensions import Self
@@ -20,6 +20,7 @@
     get_sync_engine,
 )
 from datafaker.dialects import Random
+from datafaker.theme import get_active_theme
 from datafaker.utils import T, get_property
 
 
@@ -147,6 +148,13 @@ def __init__(
         self.config: MutableMapping[str, Any] = settings.config
         self.metadata = settings.metadata
         self._table_entries: list[TableEntry] = []
+        theme = get_active_theme()
+        self._table_theme = Theme(
+            default_color=theme.data,
+            vertical_color=theme.line,
+            horizontal_color=theme.line,
+            junction_color=theme.line,
+        )
         tables_config: MutableMapping = self.config.get("tables", {})
         if not isinstance(tables_config, MutableMapping):
             tables_config = {}
@@ -195,7 +203,7 @@ def print_table(
         :param headings: List of headings for the table.
         :param rows: List of rows of values.
         """
-        output = PrettyTable()
+        output = ColorTable(theme=self._table_theme)
         output.field_names = headings
         for row in rows:
             # Hopefully PrettyTable will accept Sequence in the future, not list
@@ -208,15 +216,19 @@ def print_table_by_columns(self, columns: Mapping[str, Sequence[str]]) -> None:
 
         :param columns: Dict of column names to the values in the column.
         """
-        output = PrettyTable()
+        output = ColorTable(theme=self._table_theme)
         row_count = max(len(col) for col in columns.values())
         for field_name, data in columns.items():
             output.add_column(field_name, list(data) + [None] * (row_count - len(data)))
         print(output)
 
     def print_results(self, result: sqlalchemy.CursorResult) -> None:
         """Print the rows resulting from a database query."""
-        self.print_table(list(result.keys()), [list(row) for row in result.all()])
+        theme = get_active_theme()
+        self.print_table(
+            [f"{theme.column}{heading}" for heading in result.keys()],
+            [list(row) for row in result.all()],
+        )
 
     def ask_save(self) -> str:
         """
@@ -277,8 +289,19 @@ def report_columns(self) -> None:
             ["tables", table.name, "columns"],
             {},
         )
+        theme = get_active_theme()
         self.print_table(
-            ["name", "type", "primary", "nullable", "foreign key", "roles"],
+            [
+                f"{theme.reset}{heading}"
+                for heading in [
+                    "name",
+                    "type",
+                    "primary",
+                    "nullable",
+                    "foreign key",
+                    "roles",
+                ]
+            ],
             [
                 [
                     name,
@@ -362,8 +385,9 @@ def do_counts(self, _arg: str) -> None:
                 return
             row_count = result.get("row_count", 0)
             self.print(self.ROW_COUNT_MSG, row_count)
+            theme = get_active_theme()
             self.print_table(
-                ["Column", "NULL count"],
+                [f"{theme.reset}Column", f"{theme.reset}NULL count"],
                 [
                     [name, row_count - count]
                     for name, count in result.items()
@@ -387,8 +411,12 @@ def do_select(self, arg: str) -> None:
                 self.print("Showing the first {} rows", max_select_rows)
             fields = list(result.keys())
             rows = result.fetchmany(max_select_rows)
+            theme = get_active_theme()
             try:
-                self.print_table(fields, rows)
+                self.print_table(
+                    [f"{theme.column}{f}" for f in fields],
+                    rows,
+                )
             except ValueError as exc:
                 self.print(self.ERROR_FAILED_DISPLAY, exc)
                 return
@@ -424,7 +452,11 @@ def do_peek(self, arg: str) -> None:
             except SQLAlchemyError as exc:
                 self.print(self.ERROR_FAILED_SQL, exc=exc, query=stmt)
                 return
-            self.print_table(list(result.keys()), result.fetchmany(max_peek_rows))
+            theme = get_active_theme()
+            self.print_table(
+                [f"{theme.column}{k}" for k in result.keys()],
+                result.fetchmany(max_peek_rows),
+            )
 
     def get_column_completions(self, text: str) -> list[str]:
         """Get completions for text to column names in the current table."""

datafaker/interactive/generators.py

@@ -12,6 +12,7 @@
 from datafaker.interactive.base import DbCmd, TableEntry, fk_column_name, or_default
 from datafaker.proposers import everything_factory
 from datafaker.proposers.base import PredefinedProposer, Proposer
+from datafaker.theme import get_active_theme
 from datafaker.utils import (
     get_columns_assigned,
     get_row_generators,
@@ -85,9 +86,12 @@ class GeneratorCmd(DbCmd):
     prompt = "(generatorconf) "
     file = None
 
-    PROPOSE_SOURCE_SAMPLE_TEXT = "Sample of actual source data: {0}..."
+    PROPOSE_SOURCE_SAMPLE_TEXT = "Sample of actual source data: {1}{0}..."
     PROPOSE_SOURCE_EMPTY_TEXT = "Source database has no data in this column."
-    PROPOSE_GENERATOR_SAMPLE_TEXT = "{index}. {name}: {fit} {sample} ..."
+    PROPOSE_GENERATOR_SAMPLE_TEXT = (
+        "{theme_reset}{index}. {theme_func}{name}:"
+        " {theme_fit}{fit} {theme_data}{sample}{theme_reset} ..."
+    )
     PRIMARY_PRIVATE_TEXT = "Primary Private"
     SECONDARY_PRIVATE_TEXT = "Secondary Private on columns {0}"
     NOT_PRIVATE_TEXT = "Not private"
@@ -266,18 +270,21 @@ def _column_metadata(self) -> list[Column]:
     def set_prompt(self) -> None:
         """Set the prompt according to the current table, column and generator."""
         (table_name, prop_info) = self._get_table_and_proposer()
+        theme = get_active_theme()
         if table_name is None:
-            self.prompt = "(generators) "
+            self.prompt = f"{theme.prompt}(generators){theme.reset} "
             return
         if prop_info is None:
-            self.prompt = f"({table_name}) "
+            self.prompt = f"{theme.prompt}({table_name}){theme.reset} "
             return
         table = self.table_metadata()
         columns = [
             c + "[pk]" if table.columns[c].primary_key else c for c in prop_info.columns
         ]
         gen = f" ({prop_info.proposer.name()})" if prop_info.proposer else ""
-        self.prompt = f"({table_name}.{','.join(columns)}{gen}) "
+        self.prompt = (
+            f"{theme.prompt}({table_name}.{','.join(columns)}{gen}){theme.reset} "
+        )
 
     def _remove_auto_src_stats(self) -> list[MutableMapping[str, Any]]:
         """
@@ -655,13 +662,16 @@ def do_compare(self, arg: str) -> None:
             ]
         }
         props: list[Proposer] = self._get_proposer_proposals()
+        theme = get_active_theme()
         table_name = self.table_name()
         for argument in args:
             if argument.isdigit():
                 n = int(argument)
                 if 0 < n <= len(props):
                     prop = props[n - 1]
-                    comparison[f"{n}. {prop.name()}"] = prop.generate_data(limit)
+                    comparison[
+                        f"{n}. {theme.function}{prop.name()}"
+                    ] = prop.generate_data(limit)
                     self._print_values_queried(table_name, n, prop)
         self.print_table_by_columns(comparison)
 
@@ -677,17 +687,22 @@ def _print_values_queried(self, table_name: str, n: int, prop: Proposer) -> None
         :param n: A number to print at the start of the output.
         :param gen: The proposer to report.
         """
+        theme = get_active_theme()
         if not prop.select_aggregate_clauses() and not prop.custom_queries():
             self.print(
-                "{0}. {1} requires no data from the source database.",
+                "{0}. {2}{1}{3} requires no data from the source database.",
                 n,
                 prop.name(),
+                theme.function,
+                theme.reset,
             )
         else:
             self.print(
-                "{0}. {1} requires the following data from the source database:",
+                "{0}. {2}{1}{3} requires the following data from the source database:",
                 n,
                 prop.name(),
+                theme.function,
+                theme.reset,
             )
             self._print_select_aggregate_query(table_name, prop)
             self._print_custom_queries(prop)
@@ -709,11 +724,15 @@ def _print_custom_queries(self, prop: Proposer) -> None:
             nominal,
             actual,
         )
+        theme = get_active_theme()
         for cq_key, cq in cqs.items():
             self.print(
-                "{0}; providing the following values: {1}",
+                "{2}{0}{3}; providing the following values: {4}{1}",
                 cq["query"],
                 cq_key2args[cq_key],
+                theme.query,
+                theme.reset,
+                theme.data,
             )
 
     def _get_custom_queries_from(
@@ -775,7 +794,15 @@ def _print_select_aggregate_query(self, table_name: str, prop: Proposer) -> None
                     n,
                 )
         select_q = get_aggregate_query([prop], table_name, self.engine)
-        self.print("{0}; providing the following values: {1}", select_q, vals)
+        theme = get_active_theme()
+        self.print(
+            "{2}{0}{3}; providing the following values: {4}{1}",
+            select_q,
+            vals,
+            theme.query,
+            theme.reset,
+            theme.data,
+        )
 
     def _get_column_data(
         self, count: int, to_str: Callable[[Any], str] = repr
@@ -802,12 +829,13 @@ def do_propose(self, _arg: str) -> None:
         The results can be compared (against a sample of the real data in
         the column and against each other) with the 'compare' command.
         """
+        theme = get_active_theme()
         limit = 5
         props = self._get_proposer_proposals()
         sample = self._get_column_data(limit)
         if sample:
             rep = [x[0] if len(x) == 1 else ",".join(x) for x in sample]
-            self.print(self.PROPOSE_SOURCE_SAMPLE_TEXT, "; ".join(rep))
+            self.print(self.PROPOSE_SOURCE_SAMPLE_TEXT, "; ".join(rep), theme.data)
         else:
             self.print(self.PROPOSE_SOURCE_EMPTY_TEXT)
         if not props:
@@ -826,6 +854,10 @@ def do_propose(self, _arg: str) -> None:
                 name=prop.name(),
                 fit=fit_s,
                 sample="; ".join(map(repr, prop.generate_data(limit))),
+                theme_func=theme.function,
+                theme_fit=theme.query,
+                theme_data=theme.data,
+                theme_reset=theme.reset,
             )
 
     def do_p(self, arg: str) -> None:

datafaker/main.py

@@ -40,6 +40,7 @@
     get_source_dsn,
     get_source_schema,
 )
+from datafaker.theme import ThemeEntry, set_active_theme
 from datafaker.utils import (
     CONFIG_SCHEMA_PATH,
     conf_logger,
@@ -148,10 +149,12 @@ def load_metadata_for_output(
 
 @app.callback()
 def main(
-    verbose: bool = Option(False, "--verbose", "-v", help="Print more information.")
+    verbose: bool = Option(False, "--verbose", "-v", help="Print more information."),
+    theme: ThemeEntry = Option(ThemeEntry.DARK, help="The colour scheme"),
 ) -> None:
     """Set the global parameters."""
     conf_logger(verbose)
+    set_active_theme(theme)
 
 
 @app.command(rich_help_panel="Configure and Extract")

datafaker/theme.py

@@ -0,0 +1,65 @@
+"""Entrypoint for the datafaker package."""
+from dataclasses import dataclass
+from enum import Enum
+
+import colorama
+
+colorama.just_fix_windows_console()
+
+
+@dataclass
+class Theme:
+    """A colour theme for DataFaker terminal output."""
+
+    prompt: str
+    column: str
+    data: str
+    function: str
+    query: str
+    line: str
+    reset: str
+
+
+class ThemeEntry(str, Enum):
+    """Themes available in the ``--theme`` option."""
+
+    NONE = "none"
+    DARK = "dark"
+    LIGHT = "light"
+
+
+THEME: dict[str, Theme] = {
+    ThemeEntry.NONE: Theme("", "", "", "", "", "", ""),
+    ThemeEntry.DARK: Theme(
+        prompt=colorama.Fore.CYAN + colorama.Style.NORMAL,  # type: ignore
+        column=colorama.Fore.GREEN + colorama.Style.NORMAL,  # type: ignore
+        data=colorama.Fore.YELLOW + colorama.Style.NORMAL,  # type: ignore
+        function=colorama.Fore.MAGENTA + colorama.Style.NORMAL,  # type: ignore
+        query=colorama.Fore.GREEN + colorama.Style.NORMAL,  # type: ignore
+        line=colorama.Fore.WHITE + colorama.Style.DIM,  # type: ignore
+        reset=colorama.Style.RESET_ALL,  # type: ignore
+    ),
+    ThemeEntry.LIGHT: Theme(
+        prompt=colorama.Fore.BLUE,  # type: ignore
+        column=colorama.Fore.GREEN,  # type: ignore
+        data=colorama.Fore.BLACK,  # type: ignore
+        function=colorama.Fore.MAGENTA,  # type: ignore
+        query=colorama.Fore.MAGENTA,  # type: ignore
+        line=colorama.Fore.LIGHTBLACK_EX,  # type: ignore
+        reset=colorama.Style.RESET_ALL,  # type: ignore
+    ),
+}
+
+
+theme_active = THEME[ThemeEntry.NONE]
+
+
+def set_active_theme(te: ThemeEntry):
+    """Set the active theme by key."""
+    global theme_active  # pylint: disable=global-statement
+    theme_active = THEME[te]
+
+
+def get_active_theme() -> Theme:
+    """Get the active theme."""
+    return theme_active

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "datafaker"
-version = "0.5.0"
+version = "0.5.1"
 description = "Generates fake SQL data"
 authors = ["Tim Band <3266052+tim-band@users.noreply.github.com>"]
 license = "MIT"
@@ -39,6 +39,7 @@ prettytable = "^3.15.1"
 fastparquet = "^2024.11.0"
 duckdb-sqlalchemy = "^1.5.2.2"
 ty = "^0.0.43"
+colorama = "^0.4.6"
 
 [tool.poetry.group.dev.dependencies]
 isort = "^5.10.1"