From ae55559c32422a0fda6c5b1ad3f0523313ccab7e Mon Sep 17 00:00:00 2001
From: Brent Yi <yibrenth@gmail.com>
Date: Tue, 25 Oct 2022 22:44:44 +0800
Subject: [PATCH] Fix completion install script for bash (#830)

* Fix completion install script for bash

* Update install.py
---
 scripts/completions/install.py | 72 ++++++++++++++++++----------------
 1 file changed, 39 insertions(+), 33 deletions(-)

diff --git a/scripts/completions/install.py b/scripts/completions/install.py
index 37cfd9fbc..adfb76e7c 100755
--- a/scripts/completions/install.py
+++ b/scripts/completions/install.py
@@ -133,46 +133,52 @@ def _update_rc(
     """
 
     # Install or uninstall `source_line`.
-    source_lines = "\n".join(
-        [
-            "",
-            "# Source nerfstudio autocompletions.",
-            "# From https://unix.stackexchange.com/questions/339954/zsh-command-not-found-compinstall-compinit-compdef",
-            "autoload -Uz compinit",
-            "compinit",
-            f"source {completions_dir / 'setup'}.{shell}",
-        ]
-    )
+    header_line = "# Source nerfstudio autocompletions."
+    if shell == "zsh":
+        source_lines = "\n".join(
+            [
+                "",
+                header_line,
+                "if ! command -v compdef &> /dev/null; then",
+                "    autoload -Uz compinit",
+                "    compinit",
+                "fi",
+                f"source {completions_dir / 'setup.zsh'}",
+            ]
+        )
+    elif shell == "bash":
+        source_lines = "\n".join(
+            [
+                "",
+                header_line,
+                f"source {completions_dir / 'setup.bash'}",
+            ]
+        )
+    else:
+        assert_never(shell)
+
     rc_path = pathlib.Path(os.environ["HOME"]) / f".{shell}rc"
+
+    # Always try to uninstall previous completions.
+    rc_source = rc_path.read_text()
+    while header_line in rc_source:
+        before_install, _, after_install = rc_source.partition(header_line)
+        source_file, _, after_install = after_install.partition("\nsource ")[2].partition("\n")
+        assert source_file.endswith(f"/completions/setup.{shell}")
+        rc_source = before_install + after_install
+        rc_path.write_text(rc_source)
+        CONSOLE.log(f":broom: Existing completions uninstalled from {rc_path}.")
+
+    # Install completions.
     if mode == "install":
-        if source_lines in rc_path.read_text():
-            CONSOLE.log(f":call_me_hand: Completions are already installed in {rc_path}. {_exclamation()}")
-            return
-
-        if not Confirm.ask(f"[bold yellow]Install to {rc_path}?", default=True):
-            CONSOLE.log(f"[bold red]Skipping install for {rc_path.name}.")
-            return
-
-        rc_path.write_text(rc_path.read_text() + source_lines)
+        assert source_lines not in rc_source
+        rc_path.write_text(rc_source.rstrip() + "\n" + source_lines)
         CONSOLE.log(
             f":person_gesturing_ok: Completions installed to {rc_path}. {_exclamation()} Open a new shell to try them"
             " out."
         )
-
-    elif mode == "uninstall":
-        if source_lines not in rc_path.read_text():
-            CONSOLE.log(f":heavy_check_mark: No completions to uninstall from {rc_path.name}.")
-            return
-
-        if not Confirm.ask(f"[bold yellow]Uninstall from {rc_path}?", default=True):
-            CONSOLE.log(f"[dim red]Skipping uninstall for {rc_path.name}.")
-            return
-
-        rc_path.write_text(rc_path.read_text().replace(source_lines, ""))
-        CONSOLE.log(f":broom: Completions uninstalled from {rc_path}.")
-
     else:
-        assert_never(mode)
+        assert mode == "uninstall"
 
 
 def main(mode: ConfigureMode = "install") -> None: