diff --git a/click_help_colors/core.py b/click_help_colors/core.py index 1b32388..86a077d 100644 --- a/click_help_colors/core.py +++ b/click_help_colors/core.py @@ -1,9 +1,13 @@ +import re + import click from .utils import _colorize, _extend_instance class HelpColorsFormatter(click.HelpFormatter): + options_regex = re.compile(r'-{1,2}[\w\-]+') + def __init__(self, headers_color=None, options_color=None, options_custom_colors=None, *args, **kwargs): self.headers_color = headers_color @@ -11,13 +15,22 @@ def __init__(self, headers_color=None, options_color=None, self.options_custom_colors = options_custom_colors super(HelpColorsFormatter, self).__init__(*args, **kwargs) - def _pick_color(self, option_name): - opt = option_name.split()[0] - if (self.options_custom_colors and - (opt in self.options_custom_colors.keys())): - return self.options_custom_colors[opt] + def _get_opt_names(self, option_name): + opts = self.options_regex.findall(option_name) + if not opts: + return [option_name] else: - return self.options_color + # Include this for backwards compatibility + opts.append(option_name.split()[0]) + return opts + + def _pick_color(self, option_name): + opts = self._get_opt_names(option_name) + for opt in opts: + if (self.options_custom_colors and + (opt in self.options_custom_colors.keys())): + return self.options_custom_colors[opt] + return self.options_color def write_usage(self, prog, args='', prefix='Usage: '): colorized_prefix = _colorize(prefix, color=self.headers_color) diff --git a/tests/test_custom_colors.py b/tests/test_custom_colors.py index b434e19..8e87c80 100644 --- a/tests/test_custom_colors.py +++ b/tests/test_custom_colors.py @@ -1,3 +1,4 @@ +import pytest import click from click_help_colors import HelpColorsGroup, HelpColorsCommand @@ -62,3 +63,83 @@ def command2(name): ' \x1b[31mcommand1\x1b[0m', ' \x1b[32mcommand2\x1b[0m' ] + + +def test_option_color(runner): + @click.group( + cls=HelpColorsGroup, + help_headers_color='yellow', + help_options_color='green', + help_options_custom_colors={'--name': 'red'} + ) + def cli(): + pass + + @cli.command() + @click.option('--name', help='The person to greet.') + def command(name): + pass + + result = runner.invoke(cli, ['command', '--help'], color=True) + assert not result.exception + assert result.output.splitlines() == [ + '\x1b[33mUsage: \x1b[0mcli command [OPTIONS]', + '', + '\x1b[33mOptions\x1b[0m:', + ' \x1b[31m--name TEXT\x1b[0m The person to greet.', + ' \x1b[32m--help\x1b[0m Show this message and exit.' + ] + + +@pytest.mark.parametrize('option_name', ['-n', '--name', '-n,']) +def test_multi_name_option_color(runner, option_name): + @click.group( + cls=HelpColorsGroup, + help_headers_color='yellow', + help_options_color='green', + help_options_custom_colors={option_name: 'red'} + ) + def cli(): + pass + + @cli.command() + @click.option('-n', '--name', help='The person to greet.') + def command(name): + pass + + result = runner.invoke(cli, ['command', '--help'], color=True) + assert not result.exception + assert result.output.splitlines() == [ + '\x1b[33mUsage: \x1b[0mcli command [OPTIONS]', + '', + '\x1b[33mOptions\x1b[0m:', + ' \x1b[31m-n, --name TEXT\x1b[0m The person to greet.', + ' \x1b[32m--help\x1b[0m Show this message and exit.' + ] + + +@pytest.mark.parametrize('option_name', ['--shout', '--no-shout']) +def test_flag_option_color(runner, option_name): + @click.group( + cls=HelpColorsGroup, + help_headers_color='yellow', + help_options_color='green', + help_options_custom_colors={option_name: 'red'} + ) + def cli(): + pass + + @cli.command() + @click.option('--shout/--no-shout', default=False) + def command(name): + pass + + result = runner.invoke(cli, ['command', '--help'], color=True) + assert not result.exception + assert result.output.splitlines() == [ + '\x1b[33mUsage: \x1b[0mcli command [OPTIONS]', + '', + '\x1b[33mOptions\x1b[0m:', + ' \x1b[31m--shout / --no-shout\x1b[0m', + ' \x1b[32m--help\x1b[0m Show this message and exit.' + ]