From a6d410a290484d15e5382d31469296aa2da8aa0e Mon Sep 17 00:00:00 2001
From: "bazel.build machine account"
 <15028808+bazel-io@users.noreply.github.com>
Date: Wed, 31 Jan 2024 11:37:07 +0100
Subject: [PATCH] [7.1.0] Add bash completion for external targets (#21149)

Completion for both repository names and targets in (fetched) external
repositories is provided based on `bazel mod dump_repo_mapping`.

RELNOTES: Bazel's Bash completion can now complete external repository
labels when using `--enable_bzlmod`.

Closes #20693.

Commit
https://github.com/bazelbuild/bazel/commit/a4c3d899212bb216abe3789a13657e7413a3eefc

PiperOrigin-RevId: 602923007
Change-Id: I7c7dfe3d90a2921f473c2d2fedbc054da67932ac

Co-authored-by: Fabian Meumertzheim <fabian@meumertzhe.im>
---
 scripts/BUILD                                 |   5 +-
 scripts/bazel-complete-template.bash          | 133 +++++++-
 src/test/py/bazel/BUILD                       |  16 +
 .../bzlmod/external_repo_completion_test.py   | 285 ++++++++++++++++++
 src/test/py/bazel/bzlmod/test_utils.py        |   1 +
 5 files changed, 435 insertions(+), 5 deletions(-)
 create mode 100644 src/test/py/bazel/bzlmod/external_repo_completion_test.py

diff --git a/scripts/BUILD b/scripts/BUILD
index badb397ca7d154..7e51e493b78c6e 100644
--- a/scripts/BUILD
+++ b/scripts/BUILD
@@ -22,7 +22,10 @@ genrule(
         ":generate_bash_completion.sh",
         "//src:bazel",
     ],
-    visibility = ["//scripts/packages:__subpackages__"],
+    visibility = [
+        "//scripts/packages:__subpackages__",
+        "//src/test/py/bazel:__pkg__",
+    ],
 )
 
 sh_test(
diff --git a/scripts/bazel-complete-template.bash b/scripts/bazel-complete-template.bash
index 2b3402c2fd9bb0..f36fb1e2afe43e 100644
--- a/scripts/bazel-complete-template.bash
+++ b/scripts/bazel-complete-template.bash
@@ -98,9 +98,8 @@ _bazel__get_workspace_path() {
   echo $workspace
 }
 
-
 # Find the current piece of the line to complete, but only do word breaks at
-# certain characters. In particular, ignore these: "':=
+# certain characters. In particular, ignore these: "':=@
 # This method also takes into account the current cursor position.
 #
 # Works with both bash 3 and 4! Bash 3 and 4 perform different word breaks when
@@ -109,13 +108,14 @@ _bazel__get_workspace_path() {
 _bazel__get_cword() {
   local cur=${COMP_LINE:0:$COMP_POINT}
   # This expression finds the last word break character, as defined in the
-  # COMP_WORDBREAKS variable, but without '=' or ':', which is not preceeded by
-  # a slash. Quote characters are also excluded.
+  # COMP_WORDBREAKS variable, but without '@', '=' or ':', which is not
+  # preceded by a slash. Quote characters are also excluded.
   local wordbreaks="$COMP_WORDBREAKS"
   wordbreaks="${wordbreaks//\'/}"
   wordbreaks="${wordbreaks//\"/}"
   wordbreaks="${wordbreaks//:/}"
   wordbreaks="${wordbreaks//=/}"
+  wordbreaks="${wordbreaks//@/}"
   local word_start=$(expr "$cur" : '.*[^\]['"${wordbreaks}"']')
   echo "${cur:$word_start}"
 }
@@ -281,6 +281,111 @@ _bazel__expand_package_name() {
   done
 }
 
+# Usage: _bazel__filter_repo_mapping <filter> <field>
+#
+# Returns all entries of the main repo's repository mapping whose apparent repo
+# name, followed by a double quote, matches the given filter. To return the
+# matching apparent names, set field to 2. To return the matching canonical
+# names, set field to 4.
+# Note: Instead of returning an empty canonical name for the main repository,
+# this function returns the string "_main" so that this case can be
+# distinguished from that of no match.
+_bazel__filter_repo_mapping() {
+  local filter=$1 field=$2
+  # 1. dump_repo_mapping '' returns a single line consisting of a minified JSON
+  #    object.
+  # 2. Transform JSON to have lines of the form "apparent_name":"canonical_name".
+  # 3. Filter by apparent repo name.
+  # 4. Replace an empty canonical name with "_main".
+  # 5. Cut out either the apparent or canonical name.
+  ${BAZEL} mod dump_repo_mapping '' --noshow_progress 2>/dev/null |
+    tr '{},' '\n' |
+    "grep" "^\"${filter}" |
+    sed 's|:""$|:"_main"|' |
+    cut -d'"' -f${field}
+}
+
+# Usage: _bazel__expand_repo_name <current>
+#
+# Returns completions for apparent repository names. Each line is of the form
+# @apparent_name or @apparent_name//, where apparent_name starts with current.
+_bazel__expand_repo_name() {
+  local current=$1
+  # If current exactly matches a repo name, also provide the @current//
+  # completion so that users can tab through to package completion, but also
+  # complete just the shorthand for "@repo_name//:repo_name".
+  _bazel__filter_repo_mapping "${current#@}" 2 |
+    sed 's|^|@|' |
+    sed "s|^${current}\$|${current} ${current}//|"
+}
+
+# Usage: _bazel__repo_root <workspace> <repo>
+#
+# Returns the absolute path to the root of the repository identified by the
+# repository part <repo> of a label. <repo> can be either of the form
+# "@apparent_name" or "@@canonical_name" and may also refer to the main
+# repository.
+_bazel__repo_root() {
+  local workspace=$1 repo=$2
+  local canonical_repo
+  if [[ "$repo" == @@ ]]; then
+    # Match the sentinel value for the main repository used by
+    # _bazel__filter_repo_mapping.
+    canonical_repo=_main
+  elif [[ "$repo" =~ ^@@ ]]; then
+    # Canonical repo names should not go through repo mapping.
+    canonical_repo=${repo#@@}
+  else
+    canonical_repo=$(_bazel__filter_repo_mapping "${repo#@}\"" 4)
+  fi
+  if [ -z "$canonical_repo" ]; then
+    return
+  fi
+  if [ "$canonical_repo" == "_main" ]; then
+    echo "$workspace"
+    return
+  fi
+  local output_base="$(${BAZEL} info output_base --noshow_progress 2>/dev/null)"
+  if [ -z "$output_base" ]; then
+    return
+  fi
+  local repo_root="$output_base/external/$canonical_repo"
+  echo "$repo_root"
+}
+
+# Usage: _bazel__expand_package_name <workspace> <current> <label-type>
+#
+# Expands packages under the potentially external repository pointed to by
+# <current>, which is expected to start with "@repo//".
+_bazel__expand_external_package_name() {
+  local workspace=$1 current=$2 label_syntax=$3
+  local repo=$(echo "$current" | cut -f1 -d/)
+  local package=$(echo "$current" | cut -f3- -d/)
+  local repo_root=$(_bazel__repo_root "$workspace" "$repo")
+  if [ -z "$repo_root" ]; then
+    return
+  fi
+  _bazel__expand_package_name "$repo_root" "" "$package" "$label_syntax" |
+    sed "s|^|${repo}//|"
+}
+
+# Usage: _bazel__expand_rules_in_external_package <workspace> <current>
+#                                                 <label-type>
+#
+# Expands rule names in the potentially external package pointed to by
+# <current>, which is expected to start with "@repo//some/pkg:".
+_bazel__expand_rules_in_external_package() {
+  local workspace=$1 current=$2 label_syntax=$3
+  local repo=$(echo "$current" | cut -f1 -d/)
+  local package=$(echo "$current" | cut -f3- -d/ | cut -f1 -d:)
+  local name=$(echo "$current" | cut -f2 -d:)
+  local repo_root=$(_bazel__repo_root "$workspace" "$repo")
+  if [ -z "$repo_root" ]; then
+    return
+  fi
+  _bazel__expand_rules_in_package "$repo_root" "" "//$package:$name" "$label_syntax"
+}
+
 # Usage: _bazel__expand_target_pattern <workspace> <displacement>
 #                                      <word> <label-syntax>
 #
@@ -290,6 +395,26 @@ _bazel__expand_package_name() {
 _bazel__expand_target_pattern() {
   local workspace=$1 displacement=$2 current=$3 label_syntax=$4
   case "$current" in
+    @*//*:*) # Expand rule names within external repository.
+      _bazel__expand_rules_in_external_package "$workspace" "$current" "$label_syntax"
+      ;;
+    @*/*) # Expand package names within external repository.
+      # Append a second slash after the repo name before performing completion
+      # if there is no second slash already.
+      if [[ "$current" =~ ^@[^/]*/$ ]]; then
+        current="$current/"
+      fi
+      _bazel__expand_external_package_name "$workspace" "$current" "$label_syntax"
+      ;;
+    @*) # Expand external repository names.
+      # Do not expand canonical repository names: Users are not expected to
+      # compose them manually and completing them based on the contents of the
+      # external directory has a high risk of returning stale results.
+      if [[ "$current" =~ ^@@ ]]; then
+        return
+      fi
+      _bazel__expand_repo_name "$current"
+      ;;
     //*:*) # Expand rule names within package, no displacement.
       if [ "${label_syntax}" = "label-package" ]; then
         compgen -S " " -W "BUILD" "$(echo current | cut -f ':' -d2)"
diff --git a/src/test/py/bazel/BUILD b/src/test/py/bazel/BUILD
index 6d089f11d59793..0434970a6867ad 100644
--- a/src/test/py/bazel/BUILD
+++ b/src/test/py/bazel/BUILD
@@ -391,3 +391,19 @@ py_test(
         ":test_base",
     ],
 )
+
+py_test(
+    name = "external_repo_completion_test",
+    size = "large",
+    srcs = ["bzlmod/external_repo_completion_test.py"],
+    data = ["//scripts:bash_completion"],
+    tags = [
+        "no_windows",  # //scripts:bash_completion does not build on Windows
+        "requires-network",
+    ],
+    deps = [
+        ":bzlmod_test_utils",
+        ":test_base",
+        requirement("bazel-runfiles"),
+    ],
+)
diff --git a/src/test/py/bazel/bzlmod/external_repo_completion_test.py b/src/test/py/bazel/bzlmod/external_repo_completion_test.py
new file mode 100644
index 00000000000000..5a0f1380d7820c
--- /dev/null
+++ b/src/test/py/bazel/bzlmod/external_repo_completion_test.py
@@ -0,0 +1,285 @@
+# pylint: disable=g-backslash-continuation
+# Copyright 2023 The Bazel Authors. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Tests the bash completion for external repositories."""
+
+import os
+import subprocess
+import tempfile
+from absl.testing import absltest
+import runfiles
+from src.test.py.bazel import test_base
+from src.test.py.bazel.bzlmod.test_utils import BazelRegistry
+from src.test.py.bazel.bzlmod.test_utils import scratchFile
+
+
+class ExternalRepoCompletionTest(test_base.TestBase):
+  """Test class for bash completion for external."""
+
+  def setUp(self):
+    test_base.TestBase.setUp(self)
+    r = runfiles.Create()
+    self.completion_script = r.Rlocation('io_bazel/scripts/bazel-complete.bash')
+    self.bazel_binary = r.Rlocation('io_bazel/src/bazel')
+
+    self.registries_work_dir = tempfile.mkdtemp(dir=self._test_cwd)
+    self.main_registry = BazelRegistry(
+        os.path.join(self.registries_work_dir, 'main')
+    )
+    self.main_registry.setModuleBasePath('projects')
+    self.projects_dir = self.main_registry.projects
+    self.maxDiff = None  # there are some long diffs in this test
+
+    self.ScratchFile(
+        '.bazelrc',
+        [
+            # The command completion script invokes bazel with arguments we
+            # don't control, so we need to import the default test .bazelrc
+            # here.
+            'import ' + self._test_bazelrc,
+            # In ipv6 only network, this has to be enabled.
+            # 'startup --host_jvm_args=-Djava.net.preferIPv6Addresses=true',
+            'common --noenable_workspace',
+            'common --registry=' + self.main_registry.getURL(),
+            # We need to have BCR here to make sure built-in modules like
+            # bazel_tools can work.
+            'common --registry=https://bcr.bazel.build',
+            # Disable yanked version check so we are not affected BCR changes.
+            'common --allow_yanked_versions=all',
+            # Make sure Bazel CI tests pass in all environments
+            'common --charset=ascii',
+        ],
+    )
+
+    self.ScratchFile(
+        'MODULE.bazel',
+        [
+            'module(name = "my_project", version = "1.0")',
+            '',
+            'bazel_dep(name = "foo", version = "1.0", repo_name = "foo")',
+            'bazel_dep(name = "foo", version = "2.0", repo_name = "foobar")',
+            'bazel_dep(name = "ext", version = "1.0")',
+            'bazel_dep(name = "ext2", version = "1.0")',
+            'multiple_version_override(',
+            '  module_name= "foo",',
+            '  versions = ["1.0", "2.0"],',
+            ')',
+            'ext = use_extension("@ext//:ext.bzl", "ext")',
+            'use_repo(ext, myrepo="repo1")',
+            'ext2 = use_extension("@ext2//:ext.bzl", "ext")',
+            'ext2.dep(name="repo1")',
+            'use_repo(ext2, myrepo2="repo1")',
+        ],
+    )
+    self.ScratchFile('pkg/BUILD', ['cc_library(name = "my_lib")'])
+    self.main_registry.createCcModule(
+        'foo',
+        '1.0',
+        {'bar': '1.0', 'ext': '1.0'},
+        {'bar': 'bar_from_foo1'},
+        extra_module_file_contents=[
+            'my_ext = use_extension("@ext//:ext.bzl", "ext")',
+            'my_ext.dep(name="repo1")',
+            'my_ext.dep(name="repo2")',
+            'my_ext.dep(name="repo5")',
+            'use_repo(my_ext, my_repo1="repo1")',
+        ],
+    )
+    self.main_registry.createCcModule(
+        'foo',
+        '2.0',
+        {'bar': '2.0', 'ext': '1.0'},
+        {'bar': 'bar_from_foo2', 'ext': 'ext_mod'},
+        extra_module_file_contents=[
+            'my_ext = use_extension("@ext_mod//:ext.bzl", "ext")',
+            'my_ext.dep(name="repo4")',
+            'use_repo(my_ext, my_repo3="repo3", my_repo4="repo4")',
+        ],
+    )
+    self.main_registry.createCcModule('bar', '1.0', {'ext': '1.0'})
+    self.main_registry.createCcModule(
+        'bar',
+        '2.0',
+        {'ext': '1.0', 'ext2': '1.0'},
+        extra_module_file_contents=[
+            'my_ext = use_extension("@ext//:ext.bzl", "ext")',
+            'my_ext.dep(name="repo3")',
+            'use_repo(my_ext, my_repo3="repo3")',
+            'my_ext2 = use_extension("@ext2//:ext.bzl", "ext")',
+            'my_ext2.dep(name="repo3")',
+            'use_repo(my_ext2, my_repo2="repo3")',
+        ],
+    )
+
+    ext_src = [
+        'def _data_repo_impl(ctx): ctx.file("BUILD")',
+        'data_repo = repository_rule(_data_repo_impl,',
+        '  attrs={"data":attr.string()},',
+        ')',
+        'def _ext_impl(ctx):',
+        '  deps = {dep.name: 1 for mod in ctx.modules for dep in mod.tags.dep}',
+        '  for dep in deps:',
+        '    data_repo(name=dep, data="requested repo")',
+        'ext=module_extension(_ext_impl,',
+        '  tag_classes={"dep":tag_class(attrs={"name":attr.string()})},',
+        ')',
+    ]
+
+    self.main_registry.createLocalPathModule('ext', '1.0', 'ext')
+    scratchFile(
+        self.projects_dir.joinpath('ext', 'BUILD'),
+        ['cc_library(name="lib_ext")'],
+    )
+    scratchFile(
+        self.projects_dir.joinpath('ext', 'tools', 'BUILD'),
+        ['cc_binary(name="tool")'],
+    )
+    scratchFile(
+        self.projects_dir.joinpath('ext', 'tools', 'zip', 'BUILD'),
+        ['cc_binary(name="zipper")'],
+    )
+    scratchFile(self.projects_dir.joinpath('ext', 'ext.bzl'), ext_src)
+    self.main_registry.createLocalPathModule('ext2', '1.0', 'ext2')
+    scratchFile(
+        self.projects_dir.joinpath('ext2', 'BUILD'),
+        ['cc_library(name="lib_ext2")'],
+    )
+    scratchFile(self.projects_dir.joinpath('ext2', 'ext.bzl'), ext_src)
+
+  def complete(self, bazel_args):
+    """Get the bash completions for the given "bazel" command line."""
+
+    # The full command line to complete as typed by the user
+    # (may end with a space).
+    comp_line = 'bazel ' + bazel_args
+    # The index of the cursor position relative to the beginning of COMP_LINE.
+    comp_point = len(comp_line)
+    # The index of the word to be completed in COMP_LINE.
+    comp_cword = len(comp_line.split(' '))
+    script = """
+source {completion_script}
+COMP_WORDS=({comp_line})
+_bazel__complete
+echo ${{COMPREPLY[*]}}
+""".format(
+        completion_script=self.completion_script,
+        comp_line=comp_line,
+    )
+    env = os.environ.copy()
+    env.update({
+        # Have the completion script use the Bazel binary provided by the test
+        # runner.
+        'BAZEL': self.bazel_binary,
+        'COMP_LINE': comp_line,
+        'COMP_POINT': str(comp_point),
+        'COMP_CWORD': str(comp_cword),
+    })
+    p = subprocess.Popen(
+        ['bash', '-c', script],
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+        env=env,
+    )
+    stdout, _ = p.communicate()
+    return stdout.decode('utf-8').split()
+
+  def testCompletions(self):
+    # The completion requires an external repository to have been fetched to
+    # complete its contents. We use RunProgram instead of RunBazel as the latter
+    # would evaluate the test .bazelrc twice, which we explicitly import in the
+    # test workspace's .bazelrc.
+    self.RunProgram([self.bazel_binary, 'fetch', '@ext//...', '@foobar//...'])
+
+    # Apparent repo names are completed.
+    self.assertCountEqual(
+        [
+            '@',
+            '@//',
+            '@bazel_tools',
+            '@local_config_platform',
+            '@foo',
+            '@foobar',
+            '@my_project',
+            '@myrepo',
+            '@myrepo2',
+            '@ext',
+            '@ext2',
+        ],
+        self.complete('build @'),
+    )
+    self.assertCountEqual(['@foo', '@foobar'], self.complete('build @fo'))
+    self.assertCountEqual(
+        ['@foo', '@foo//', '@foobar'], self.complete('build @foo')
+    )
+    self.assertCountEqual(
+        ['@foobar', '@foobar//'], self.complete('build @foobar')
+    )
+    self.assertCountEqual(['@my_project'], self.complete('build @my_'))
+    self.assertCountEqual([], self.complete('build @does_not_exist'))
+
+    # Canonical repo names are not completed.
+    self.assertCountEqual([], self.complete('build @@'))
+    self.assertCountEqual([], self.complete('build @@foo~2.'))
+
+    # Packages are completed in external repos with apparent repo names.
+    self.assertCountEqual(
+        ['@ext//tools/', '@ext//tools:'], self.complete('build @ext//tool')
+    )
+    self.assertCountEqual(
+        ['@ext//tools/zip/', '@ext//tools/zip:'],
+        self.complete('build @ext//tools/zi'),
+    )
+    self.assertCountEqual(
+        ['@my_project//pkg/', '@my_project//pkg:'],
+        self.complete('build @my_project//p'),
+    )
+    self.assertCountEqual(['@//pkg/', '@//pkg:'], self.complete('build @//p'))
+    self.assertCountEqual([], self.complete('build @does_not_exist//'))
+
+    # Packages are completed in external repos with canonical repo names.
+    self.assertCountEqual(
+        ['@@ext~1.0//tools/', '@@ext~1.0//tools:'],
+        self.complete('build @@ext~1.0//tool'),
+    )
+    self.assertCountEqual(
+        ['@@ext~1.0//tools/zip/', '@@ext~1.0//tools/zip:'],
+        self.complete('build @@ext~1.0//tools/zi'),
+    )
+    self.assertCountEqual(
+        ['@@//pkg/', '@@//pkg:'], self.complete('build @@//p')
+    )
+    self.assertCountEqual([], self.complete('build @@does_not_exist//'))
+
+    # Targets are completed in external repos with apparent repo names.
+    self.assertCountEqual(['@foobar//:'], self.complete('build @foobar/'))
+    self.assertCountEqual(['@foobar//:'], self.complete('build @foobar//'))
+    # Completions operate on the last word, which is broken on ':'.
+    self.assertCountEqual(['lib_foo'], self.complete('build @foobar//:'))
+    self.assertCountEqual(
+        ['zipper'], self.complete('build @ext//tools/zip:zipp')
+    )
+    self.assertCountEqual(
+        ['my_lib'], self.complete('build @my_project//pkg:my_')
+    )
+
+    # Targets are completed in external repos with canonical repo names.
+    self.assertCountEqual(['lib_foo'], self.complete('build @@foo~2.0//:'))
+    self.assertCountEqual(
+        ['zipper'], self.complete('build @@ext~1.0//tools/zip:zipp')
+    )
+    self.assertCountEqual(['my_lib'], self.complete('build @@//pkg:my_'))
+
+
+if __name__ == '__main__':
+  absltest.main()
diff --git a/src/test/py/bazel/bzlmod/test_utils.py b/src/test/py/bazel/bzlmod/test_utils.py
index 78102db66ecbac..8cdea10582033c 100644
--- a/src/test/py/bazel/bzlmod/test_utils.py
+++ b/src/test/py/bazel/bzlmod/test_utils.py
@@ -48,6 +48,7 @@ def integrity(data):
 
 def scratchFile(path, lines=None):
   """Creates a file at the given path with the given content."""
+  os.makedirs(str(path.parent), exist_ok=True)
   with open(str(path), 'w') as f:
     if lines:
       for l in lines: