diff --git a/README.md b/README.md index 8313417..26d52fc 100644 --- a/README.md +++ b/README.md @@ -198,7 +198,7 @@ Quick Args options: - `suffix`: Add suffix to all values - `no_labels`: Don't show labels - `no_values`: Don't show values -- `colors`: List of color names +- `colors`: List of color names or ANSI color codes ## Background diff --git a/termgraph/chart.py b/termgraph/chart.py index 5adba1a..cc054e5 100644 --- a/termgraph/chart.py +++ b/termgraph/chart.py @@ -90,7 +90,7 @@ def _print_header(self) -> None: # Print categories' names above the graph. for i in range(len(self.data.categories)): if colors is not None and isinstance(colors, list): - sys.stdout.write(f"\033[{colors[i]}m") # Start to write colorized. + sys.stdout.write(f"\033[38;5;{colors[i]}m") # Start to write colorized. sys.stdout.write(f"{self.tick} {self.data.categories[i]} ") if colors: @@ -396,7 +396,7 @@ def draw(self) -> None: break if color: - sys.stdout.write(f"\033[{color}m") + sys.stdout.write(f"\033[38;5;{color}m") for row in result_list: print(*row) diff --git a/termgraph/constants.py b/termgraph/constants.py index bc59e0b..dda827d 100644 --- a/termgraph/constants.py +++ b/termgraph/constants.py @@ -15,13 +15,13 @@ TICK = "▇" SM_TICK = "▏" -# ANSI escape SGR Parameters color codes +# ANSI color codes for 256-color mode AVAILABLE_COLORS = { - "red": 91, - "blue": 94, - "green": 92, - "magenta": 95, - "yellow": 93, - "black": 90, - "cyan": 96, + "red": 9, + "blue": 12, + "green": 10, + "magenta": 13, + "yellow": 11, + "black": 0, + "cyan": 14, } \ No newline at end of file diff --git a/termgraph/termgraph.py b/termgraph/termgraph.py index c8898f5..2740d09 100644 --- a/termgraph/termgraph.py +++ b/termgraph/termgraph.py @@ -3,9 +3,8 @@ import sys from datetime import datetime, timedelta from colorama import just_fix_windows_console -import os -import re import importlib.metadata +from typing import NoReturn from .constants import AVAILABLE_COLORS, DAYS from .data import Data @@ -146,7 +145,7 @@ def chart(data_obj: Data, args: dict, colors: list) -> None: chart_obj.draw() -def _extract_colors(data_obj: Data, args: dict) -> list: +def _extract_colors(data_obj: Data, args: dict) -> list[int]: """Extract and validate colors from args based on data dimensions. Args: @@ -156,7 +155,7 @@ def _extract_colors(data_obj: Data, args: dict) -> list: Returns: List of color codes for each category """ - colors = [] + colors: list[int] = [] # Determine number of categories from data dimensions if data_obj.dims and len(data_obj.dims) > 1: @@ -165,34 +164,12 @@ def _extract_colors(data_obj: Data, args: dict) -> list: len_categories = 1 # If user inserts colors, they should be as many as the categories. - if args.get("color") is not None: - # Decompose arguments for Windows - if os.name == "nt": - colorargs = re.findall(r"[a-z]+", args["color"][0]) - if len(colorargs) != len_categories: - print(">> Error: Color and category array sizes don't match") - for color in colorargs: - if color not in AVAILABLE_COLORS: - print( - ">> Error: invalid color. choose from 'red', 'blue', 'green', 'magenta', 'yellow', 'black', 'cyan'" - ) - sys.exit(2) - else: - if len(args["color"]) != len_categories: - print(">> Error: Color and category array sizes don't match") - for color in args["color"]: - if color not in AVAILABLE_COLORS: - print( - ">> Error: invalid color. choose from 'red', 'blue', 'green', 'magenta', 'yellow', 'black', 'cyan'" - ) - sys.exit(2) - - if os.name == "nt": - for color in colorargs: - colors.append(AVAILABLE_COLORS.get(color)) - else: - for color in args["color"]: - colors.append(AVAILABLE_COLORS.get(color)) + raw_colors = args.get("color") + if raw_colors is not None: + if len(raw_colors) != len_categories: + print(">> Error: Color and category array sizes don't match") + + colors = [_validate_color(color) for color in raw_colors] # If user hasn't inserted colors, pick the first n colors # from the dict (n = number of categories). @@ -201,6 +178,37 @@ def _extract_colors(data_obj: Data, args: dict) -> list: return colors + +def _invalid_color_error() -> NoReturn: + """Print the standard invalid-color message and exit.""" + print( + ">> Error: invalid color. choose from 'red', 'blue', 'green', 'magenta', 'yellow', 'black', 'cyan', or a 256-color ANSI code (0-255)" + ) + sys.exit(2) + + +def _validate_color(color: str) -> int: + """Validate a single color value and return its ANSI code. + + A color is valid if it is either a name found in AVAILABLE_COLORS + (e.g. "red") or a 256-color ANSI code in the range 0-255 (passed as + a string). Anything else exits the program with + error 2. + """ + if color in AVAILABLE_COLORS: + return AVAILABLE_COLORS[color] + + try: + code = int(color) + except Exception: + _invalid_color_error() + + if 0 <= code <= 255: + return code + + _invalid_color_error() + + def read_data(args: dict) -> tuple[list, list, list, list]: """Read data from a file or stdin and returns it. @@ -224,10 +232,7 @@ def read_data(args: dict) -> tuple[list, list, list, list]: def calendar_heatmap(data: dict, labels: list, args: dict) -> None: """Print a calendar heatmap.""" - if args["color"]: - colornum = AVAILABLE_COLORS.get(args["color"][0]) - else: - colornum = AVAILABLE_COLORS.get("blue") + colornum = _validate_color(args["color"][0]) if args["color"] else AVAILABLE_COLORS["blue"] dt_dict = {} for i in range(len(labels)): @@ -311,4 +316,4 @@ def normalize(data: list, width: int) -> list: if __name__ == "__main__": - main() + main() \ No newline at end of file diff --git a/termgraph/utils.py b/termgraph/utils.py index 56203b7..383b7e6 100644 --- a/termgraph/utils.py +++ b/termgraph/utils.py @@ -65,7 +65,7 @@ def print_row_core( sys.stdout.write(SM_TICK) else: if color: - sys.stdout.write(f"\033[{color}m") # Start to write colorized. + sys.stdout.write(f"\033[38;5;{color}m") # Start to write colorized. for _ in range(num_blocks): sys.stdout.write(tick)