nixos-rebuild-ng: add shell completion via shtab

pkgs/by-name/ni/nixos-rebuild-ng/README.md

@@ -109,7 +109,7 @@ ruff format .
   `system.switch.enableNg` for `switch-to-configuration-ng`
 - [ ] Improve documentation
 - [x] `nixos-rebuild repl`
-- [ ] Generate tab completion via [`shtab`](https://docs.iterative.ai/shtab/)
+- [x] Generate tab completion via [`shtab`](https://docs.iterative.ai/shtab/)
 - [x] Reduce build closure
 
 ## TODON'T

pkgs/by-name/ni/nixos-rebuild-ng/package.nix

@@ -1,5 +1,6 @@
 {
   lib,
+  stdenv,
   installShellFiles,
   mkShell,
   nix,
@@ -7,8 +8,12 @@
   python3,
   python3Packages,
   runCommand,
+  withShellCompletion ? (stdenv.buildPlatform.canExecute stdenv.hostPlatform),
   withNgSuffix ? true,
 }:
+let
+  executable = if withNgSuffix then "nixos-rebuild-ng" else "nixos-rebuild";
+in
 python3Packages.buildPythonApplication rec {
   pname = "nixos-rebuild-ng";
   version = "0.0.0";
@@ -23,9 +28,13 @@ python3Packages.buildPythonApplication rec {
     tabulate
   ];
 
-  nativeBuildInputs = [
-    installShellFiles
-  ];
+  nativeBuildInputs =
+    [
+      installShellFiles
+    ]
+    ++ lib.optionals withShellCompletion [
+      python3Packages.shtab
+    ];
 
   propagatedBuildInputs = [
     # Make sure that we use the Nix package we depend on, not something
@@ -42,12 +51,14 @@ python3Packages.buildPythonApplication rec {
   postInstall =
     ''
       installManPage ${nixos-rebuild}/share/man/man8/nixos-rebuild.8
-
-      installShellCompletion \
-        --bash ${nixos-rebuild}/share/bash-completion/completions/_nixos-rebuild
+    ''
+    + lib.optionalString withShellCompletion ''
+      installShellCompletion --cmd ${executable} \
+        --bash <(shtab --shell bash nixos_rebuild.get_main_parser) \
+        --zsh <(shtab --shell zsh nixos_rebuild.get_main_parser)
     ''
     + lib.optionalString withNgSuffix ''
-      mv $out/bin/nixos-rebuild $out/bin/nixos-rebuild-ng
+      mv $out/bin/nixos-rebuild $out/bin/${executable}
     '';
 
   nativeCheckInputs = with python3Packages; [
@@ -98,6 +109,6 @@ python3Packages.buildPythonApplication rec {
     homepage = "https://github.com/NixOS/nixpkgs/tree/master/pkgs/by-name/ni/nixos-rebuild-ng";
     license = lib.licenses.mit;
     maintainers = [ lib.maintainers.thiagokokada ];
-    mainProgram = if withNgSuffix then "nixos-rebuild-ng" else "nixos-rebuild";
+    mainProgram = executable;
   };
 }

pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/__init__.py

@@ -73,32 +73,101 @@ def get_parser() -> tuple[argparse.ArgumentParser, dict[str, argparse.ArgumentPa
         add_help=False,
         allow_abbrev=False,
     )
-    main_parser.add_argument("--help", "-h", action="store_true")
-    main_parser.add_argument("--file", "-f")
-    main_parser.add_argument("--attr", "-A")
-    main_parser.add_argument("--flake", nargs="?", const=True)
-    main_parser.add_argument("--no-flake", dest="flake", action="store_false")
-    main_parser.add_argument("--install-bootloader", action="store_true")
-    main_parser.add_argument("--install-grub", action="store_true")  # deprecated
-    main_parser.add_argument("--profile-name", "-p", default="system")
-    main_parser.add_argument("--specialisation", "-c")
-    main_parser.add_argument("--rollback", action="store_true")
-    main_parser.add_argument("--upgrade", action="store_true")
-    main_parser.add_argument("--upgrade-all", action="store_true")
-    main_parser.add_argument("--json", action="store_true")
-    main_parser.add_argument("--sudo", action="store_true")
-    main_parser.add_argument("--ask-sudo-password", action="store_true")
-    main_parser.add_argument("--use-remote-sudo", action="store_true")  # deprecated
-    main_parser.add_argument("--no-ssh-tty", action="store_true")  # deprecated
-    main_parser.add_argument("--fast", action="store_true")
-    main_parser.add_argument("--build-host")
-    main_parser.add_argument("--target-host")
-    main_parser.add_argument("--no-build-nix", action="store_true")  # deprecated
+    main_parser.add_argument("--help", "-h", action="store_true", help="Show manpage")
+    main_parser.add_argument(
+        "--file", "-f", help="Enable and build the NixOS system from the specified file"
+    )
+    main_parser.add_argument(
+        "--attr",
+        "-A",
+        help="Enable and build the NixOS system from nix file and use the "
+        + "specified attribute path from file specified by the --file option",
+    )
+    main_parser.add_argument(
+        "--flake",
+        nargs="?",
+        const=True,
+        help="Build the NixOS system from the specified flake",
+    )
+    main_parser.add_argument(
+        "--no-flake",
+        dest="flake",
+        action="store_false",
+        help="Do not imply --flake if /etc/nixos/flake.nix exists",
+    )
+    main_parser.add_argument(
+        "--install-bootloader",
+        action="store_true",
+        help="Causes the boot loader to be (re)installed on the device specified "
+        + "by the relevant configuration options",
+    )
+    main_parser.add_argument(
+        "--install-grub",
+        action="store_true",
+        help="Deprecated, use '--install-bootloader' instead",
+    )
+    main_parser.add_argument(
+        "--profile-name",
+        "-p",
+        default="system",
+        help="Use nix profile /nix/var/nix/profiles/system-profiles/<profile-name>",
+    )
+    main_parser.add_argument(
+        "--specialisation", "-c", help="Activates given specialisation"
+    )
+    main_parser.add_argument(
+        "--rollback",
+        action="store_true",
+        help="Roll back to the previous configuration",
+    )
+    main_parser.add_argument(
+        "--upgrade",
+        action="store_true",
+        help="Update the root user's channel named 'nixos' before rebuilding "
+        + "the system and channels which have a file named '.update-on-nixos-rebuild'",
+    )
+    main_parser.add_argument(
+        "--upgrade-all",
+        action="store_true",
+        help="Same as --upgrade, but updates all root user's channels",
+    )
+    main_parser.add_argument(
+        "--json",
+        action="store_true",
+        help="JSON output, only implemented for 'list-generations' right now",
+    )
+    main_parser.add_argument(
+        "--ask-sudo-password",
+        action="store_true",
+        help="Asks for sudo password for remote activation, implies --sudo",
+    )
+    main_parser.add_argument(
+        "--sudo", action="store_true", help="Prefixes activation commands with sudo"
+    )
+    main_parser.add_argument(
+        "--use-remote-sudo",
+        action="store_true",
+        help="Deprecated, use '--sudo' instead",
+    )
+    main_parser.add_argument("--no-ssh-tty", action="store_true", help="Deprecated")
+    main_parser.add_argument(
+        "--fast",
+        action="store_true",
+        help="Skip possibly expensive operations",
+    )
+    main_parser.add_argument("--build-host", help="Specifies host to perform the build")
+    main_parser.add_argument("--target-host", help="Specifies the NixOS target host")
+    main_parser.add_argument("--no-build-nix", action="store_true", help="Deprecated")
     main_parser.add_argument("action", choices=Action.values(), nargs="?")
 
     return main_parser, sub_parsers
 
 
+# For shtab to generate completions
+def get_main_parser() -> argparse.ArgumentParser:
+    return get_parser()[0]
+
+
 def parse_args(
     argv: list[str],
 ) -> tuple[argparse.Namespace, dict[str, argparse.Namespace]]: