pw_ide: Add command to install prototype extension

Change-Id: I23cc6c612f97be01a6a58555222b1e9cb5176224
Reviewed-on: https://pigweed-review.googlesource.com/c/pigweed/pigweed/+/162412
Commit-Queue: Chad Norvell <chadnorvell@google.com>
Reviewed-by: Anthony DiGirolamo <tonymd@google.com>

3 files changed, 155 insertions, 4 deletions

diff --git a/pw_ide/py/pw_ide/cli.py b/pw_ide/py/pw_ide/cli.py
index 8e0c21655..ebd5ce64c 100644
--- a/pw_ide/py/pw_ide/cli.py
+++ b/pw_ide/py/pw_ide/cli.py
@@ -350,6 +350,12 @@ def _build_argument_parser() -> argparse.ArgumentParser:
         metavar='SETTINGS_TYPE',
         help='do not update these settings types',
     )
+    parser_vscode.add_argument(
+        '--install-extension',
+        dest='should_install_extension',
+        action='store_true',
+        help='install the experimental extension',
+    )
 
     return parser_root
 
diff --git a/pw_ide/py/pw_ide/commands.py b/pw_ide/py/pw_ide/commands.py
index 6263f022a..399dd332a 100644
--- a/pw_ide/py/pw_ide/commands.py
+++ b/pw_ide/py/pw_ide/commands.py
@@ -49,7 +49,11 @@ from pw_ide.settings import (
 from pw_ide.status_reporter import StatusReporter
 
 from pw_ide import vscode
-from pw_ide.vscode import VscSettingsManager, VscSettingsType
+from pw_ide.vscode import (
+    install_extension_from_vsix,
+    VscSettingsManager,
+    VscSettingsType,
+)
 
 _LOG = logging.getLogger(__package__)
 env = pigweed_environment()
@@ -118,6 +122,7 @@ def cmd_setup(
 def cmd_vscode(
     include: Optional[List[VscSettingsType]] = None,
     exclude: Optional[List[VscSettingsType]] = None,
+    should_install_extension: bool = False,
     reporter: StatusReporter = StatusReporter(),
     pw_ide_settings: PigweedIdeSettings = PigweedIdeSettings(),
 ) -> None:
@@ -242,6 +247,18 @@ def cmd_vscode(
                 f'{verb} Visual Studio Code active ' f'{settings_type.value}'
             )
 
+    if should_install_extension:
+        reporter.new("Installing Visual Studio Code extension")
+
+        try:
+            install_extension_from_vsix(reporter)
+        except FileNotFoundError:
+            reporter.err("Could not find Visual Studio Code")
+            sys.exit(1)
+        except subprocess.CalledProcessError:
+            reporter.err("Failed to install extension!")
+            sys.exit(1)
+
 
 def _process_compdbs(  # pylint: disable=too-many-locals
     reporter: StatusReporter,
diff --git a/pw_ide/py/pw_ide/vscode.py b/pw_ide/py/pw_ide/vscode.py
index 25dbb199e..fe7b1ede4 100644
--- a/pw_ide/py/pw_ide/vscode.py
+++ b/pw_ide/py/pw_ide/vscode.py
@@ -65,8 +65,12 @@ import json
 import os
 from pathlib import Path
 import platform
+import shutil
+import subprocess
 from typing import Any, Dict, List, Optional, OrderedDict
 
+from pw_cli.env import pigweed_environment
+
 from pw_ide.activate import BashShellModifier
 from pw_ide.cpp import ClangdSettings, CppIdeFeaturesState
 
@@ -79,6 +83,11 @@ from pw_ide.editors import (
 
 from pw_ide.python import PythonPaths
 from pw_ide.settings import PigweedIdeSettings
+from pw_ide.status_reporter import StatusReporter
+
+env = pigweed_environment()
+
+_VSIX_DIR = Path(env.PW_ROOT) / 'pw_ide' / 'vscode'
 
 
 def _vsc_os(system: str = platform.system()):
@@ -99,15 +108,15 @@ def _activated_env() -> OrderedDict[str, Any]:
     # Not all environments have an actions.json, which this ultimately relies
     # on (e.g. tests in CI). No problem, just return an empty dict instead.
     try:
-        env = (
+        activated_env = (
             BashShellModifier(env_only=True, path_var='${env:PATH}')
             .modify_env()
             .env_mod
         )
     except (FileNotFoundError, json.JSONDecodeError):
-        env = dict()
+        activated_env = dict()
 
-    return OrderedDict(env)
+    return OrderedDict(activated_env)
 
 
 def _local_terminal_integrated_env() -> Dict[str, Any]:
@@ -424,3 +433,122 @@ class VscSettingsManager(EditorSettingsManager[VscSettingsType]):
         VscSettingsType.EXTENSIONS: _default_extensions,
         VscSettingsType.LAUNCH: _default_launch,
     }
+
+
+def _prompt_for_path(reporter: StatusReporter) -> Path:
+    reporter.info(
+        [
+            "Hmmm... I can't seem to find your Visual Studio Code binary path!",
+            "You can provide it manually here, or Ctrl-C to cancel.",
+        ]
+    )
+
+    path = Path(input("> "))
+
+    if path.exists():
+        return path
+
+    reporter.err("Nothing found there!")
+    raise FileNotFoundError()
+
+
+# TODO(chadnorvell): Replace this when we support Python 3.11 with:
+# _PathData = Tuple[Optional[str], *Tuple[str]]
+_PathData = List[Optional[str]]
+
+
+def _try_code_path(path_list: _PathData) -> Optional[Path]:
+    root, *rest = path_list
+
+    if root is None:
+        return None
+
+    path = Path(root)
+
+    for part in rest:
+        if part is None:
+            return None
+
+        path /= part
+
+    return path
+
+
+def _try_each_code_path(
+    reporter: StatusReporter, *path_lists: _PathData
+) -> Path:
+    for path_list in path_lists:
+        if (path := _try_code_path(path_list)) is not None:
+            return path
+
+    if (path_str := shutil.which('code')) is not None:
+        return Path(path_str)
+
+    return _prompt_for_path(reporter)
+
+
+def _get_vscode_exe_path(
+    reporter: StatusReporter, system: str = platform.system()
+) -> Path:
+    if system == 'Darwin':
+        return _try_each_code_path(
+            reporter,
+            [
+                '/Applications',
+                'Visual Studio Code.app',
+                'Contents',
+                'Resources',
+                'app',
+                'bin',
+                'code',
+            ],
+        )
+
+    if system == 'Windows':
+        return _try_each_code_path(
+            reporter,
+            [
+                os.getenv('APPDATA'),
+                'Local',
+                'Programs',
+                'Microsoft VS Code',
+                'bin',
+                'code.exe',
+            ],
+            [
+                os.getenv('LOCALAPPDATA'),
+                'Local',
+                'Programs',
+                'Microsoft VS Code',
+                'bin',
+                'code.exe',
+            ],
+        )
+
+    if system == 'Linux':
+        return _try_each_code_path(
+            reporter,
+            ['/usr', 'bin', 'code'],
+            ['/usr', 'local', 'bin', 'code'],
+        )
+
+    return _prompt_for_path(reporter)
+
+
+def _get_latest_extension_vsix() -> Path:
+    return sorted(_VSIX_DIR.glob('*.vsix'), reverse=True)[0]
+
+
+def install_extension_from_vsix(reporter: StatusReporter) -> None:
+    """Install the latest pre-built VSC extension from its VSIX file.
+
+    Normally, extensions are installed from the VSC extension marketplace.
+    This will instead install the extension directly from a file.
+    """
+    extension_path = _get_latest_extension_vsix()
+    vscode_exe_path = _get_vscode_exe_path(reporter)
+
+    reporter.info(f"Path: {vscode_exe_path}")
+    subprocess.run(
+        [vscode_exe_path, '--install-extension', extension_path], check=True
+    )