Fix issue with options with multiple names

click_help_colors/core.py

@@ -1,23 +1,36 @@
+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
         self.options_color = options_color
         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)

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.'
+    ]