From d17d5341ceb48d7e4f0b7fe8260c947f30edc298 Mon Sep 17 00:00:00 2001
From: binaryDiv <binarydiv@gmail.com>
Date: Sat, 12 Jun 2021 02:02:32 +0200
Subject: [PATCH] eepprog.py: Implement Click CategorizedGroup

---
 client/eepprog.py       |  9 +++---
 client/helpers/click.py | 69 +++++++++++++++++++++++++++++++++++++++++
 2 files changed, 74 insertions(+), 4 deletions(-)
 create mode 100644 client/helpers/click.py

diff --git a/client/eepprog.py b/client/eepprog.py
index 617fe3f..2a9641f 100755
--- a/client/eepprog.py
+++ b/client/eepprog.py
@@ -4,9 +4,10 @@ import click
 import serial.tools.list_ports
 
 from helpers import CliContext, Logger, EepromProgrammer
+from helpers.click import CategorizedGroup
 
 
-@click.group()
+@click.group(cls=CategorizedGroup)
 @click.option('--device', '-d', default='/dev/ttyUSB0', show_default=True,
               metavar='DEVICE', help="Set the serial device")
 @click.option('--baud', '-b', default=38400, show_default=True,
@@ -29,7 +30,7 @@ def eepprog(ctx: click.Context, device: str, baud: int, verbose: bool) -> None:
     ctx.call_on_close(eeprom_programmer.close)
 
 
-@eepprog.command('hello', short_help='Say hello. :)')
+@eepprog.command('hello', category='Test commands', short_help='Say hello. :)')
 @click.pass_obj
 def hello(context: CliContext) -> None:
     """
@@ -46,7 +47,7 @@ def hello(context: CliContext) -> None:
     context.logger.success('yay!')
 
 
-@eepprog.command('test', short_help="Send INIT command to the programmer and read answer")
+@eepprog.command('test', category='Test commands', short_help="Send INIT command to the programmer and read answer")
 @click.pass_obj
 def test(context: CliContext) -> None:
     """
@@ -57,7 +58,7 @@ def test(context: CliContext) -> None:
     context.eeprom_programmer.test_command()
 
 
-@eepprog.command('list-devices', short_help="List available serial ports")
+@eepprog.command('list-devices', category='Helper commands', short_help="List available serial ports")
 @click.pass_obj
 def list_devices(context: CliContext) -> None:
     """
diff --git a/client/helpers/click.py b/client/helpers/click.py
new file mode 100644
index 0000000..ff38a6a
--- /dev/null
+++ b/client/helpers/click.py
@@ -0,0 +1,69 @@
+import click
+import collections
+
+
+class CategorizedGroup(click.Group):
+    """
+    Click command group that categorizes subcommands in the help text by the additional 'category' attribute.
+    Commands that don't have a 'category' attribute are grouped under the default category "Commands" as usual.
+
+    Example usage:
+    `@categorized_group.command('foobar', category='Debug commands')`
+    """
+    default_category: str
+
+    def __init__(self, name=None, commands=None, default_category: str = 'Commands', **attrs):
+        super().__init__(name, commands, **attrs)
+        self.commands = commands or collections.OrderedDict()
+        self.default_category = default_category
+
+    def list_commands(self, ctx):
+        """ List commands in the order they were added to the group. """
+        return self.commands
+
+    def command(self, *args, **kwargs):
+        """ Extends the command decorator by setting a category attribute on the command. """
+        category = kwargs.pop('category', None)
+        orig_decorator = super().command(*args, **kwargs)
+
+        def decorator(f):
+            cmd = orig_decorator(f)
+            if category:
+                cmd.category = category
+            return cmd
+
+        return decorator
+
+    def format_commands(self, ctx, formatter):
+        """
+        Prints commands in help messages.
+        This version groups them into categories and is based on the original Group.format_commands method.
+        """
+        categorized_commands = {self.default_category: []}
+
+        # Collect subcommands grouped by category
+        for subcommand in self.list_commands(ctx):
+            cmd = self.get_command(ctx, subcommand)
+            if cmd is None or cmd.hidden:
+                continue
+
+            category = getattr(cmd, 'category', self.default_category)
+            if categorized_commands.get(category) is None:
+                categorized_commands[category] = []
+
+            categorized_commands[category].append((subcommand, cmd))
+
+        # Print help sections for each subcommand category
+        for category, subcommands in categorized_commands.items():
+            if not len(subcommands):
+                continue
+
+            limit = formatter.width - 6 - max(len(cmd[0]) for cmd in subcommands)
+
+            rows = []
+            for subcommand, cmd in subcommands:
+                help_str = cmd.get_short_help_str(limit)
+                rows.append((subcommand, help_str))
+
+            with formatter.section(category):
+                formatter.write_dl(rows)