Add a coverage_tool attribute to py_runtime target.

If set, and coverage is enabled, the path to this target is used
instead of the PYTHON_COVERAGE environment variable.  This permits
use of in-build versions of the coverage tool, closing bazelbuild#14436.

src/main/java/com/google/devtools/build/lib/bazel/rules/python/BazelPythonSemantics.java

@@ -92,12 +92,18 @@ public boolean prohibitHyphensInPackagePaths() {
   public void collectRunfilesForBinary(
       RuleContext ruleContext, Runfiles.Builder builder, PyCommon common, CcInfo ccInfo) {
     addRuntime(ruleContext, common, builder);
+    if (ruleContext.getConfiguration().isCodeCoverageEnabled()) {
+      mainClass = addCoverageSupport(ruleContext, common, builder);
+    }
   }
 
   @Override
   public void collectDefaultRunfilesForBinary(
       RuleContext ruleContext, PyCommon common, Runfiles.Builder builder) {
     addRuntime(ruleContext, common, builder);
+    if (ruleContext.getConfiguration().isCodeCoverageEnabled()) {
+      mainClass = addCoverageSupport(ruleContext, common, builder);
+    }
   }
 
   @Override
@@ -154,6 +160,9 @@ private static void createStubFile(
     // first-stage.
     String pythonBinary = getPythonBinary(ruleContext, common, bazelConfig);
 
+    // The python code coverage tool to use, if any.
+    String coverageTool = getCoverageTool(ruleContext, common, bazelConfig);
+
     // Version information for host config diagnostic warning.
     PythonVersion attrVersion = PyCommon.readPythonVersionFromAttribute(ruleContext.attributes());
     boolean attrVersionSpecifiedExplicitly = attrVersion != null;
@@ -172,6 +181,7 @@ private static void createStubFile(
                 Substitution.of(
                     "%main%", common.determineMainExecutableSource(/*withWorkspaceName=*/ true)),
                 Substitution.of("%python_binary%", pythonBinary),
+                Substitution.of("%coverage_tool%", coverageTool == null ? "" : coverageTool),
                 Substitution.of("%imports%", Joiner.on(":").join(common.getImports().toList())),
                 Substitution.of("%workspace_name%", ruleContext.getWorkspaceName()),
                 Substitution.of("%is_zipfile%", boolToLiteral(isForZipFile)),
@@ -463,6 +473,31 @@ private static String getPythonBinary(
     return pythonBinary;
   }
 
+  private static void addCoverageSupport(
+      RuleContext ruleContext, PyCommon common, Runfiles.Builder builder) {
+    PyRuntimeInfo provider = getRuntime(ruleContext, common);
+    if (provider != null && provider.getCoverageTool() != null) {
+      builder.addArtifact(provider.getCoverageTool());
+    }
+  }
+
+  @Nullable
+  private static String getCoverageTool(
+      RuleContext ruleContext, PyCommon common, BazelPythonConfiguration bazelConfig) {
+    if (!ruleContext.getConfiguration().isCodeCoverageEnabled()) {
+      return null;
+    }
+    String coverageTool = null;
+    PyRuntimeInfo provider = getRuntime(ruleContext, common);
+    if (provider != null && provider.getCoverageTool() != null) {
+      PathFragment workspaceName =
+          PathFragment.create(ruleContext.getRule().getPackage().getWorkspaceName());
+      coverageTool =
+          workspaceName.getRelative(provider.getCoverageTool().getRunfilesPath()).getPathString();
+    }
+    return coverageTool;
+  }
+
   private static String getStubShebang(RuleContext ruleContext, PyCommon common) {
     PyRuntimeInfo provider = getRuntime(ruleContext, common);
     if (provider != null) {

src/main/java/com/google/devtools/build/lib/bazel/rules/python/python_stub_template.txt

@@ -332,7 +332,7 @@ def Main():
   if python_program is None:
     raise AssertionError('Could not find python binary: ' + PYTHON_BINARY)
 
-  cov_tool = os.environ.get('PYTHON_COVERAGE')
+  cov_tool = '%coverage_tool%' or os.environ.get('PYTHON_COVERAGE')
   if cov_tool:
     # Inhibit infinite recursion:
     del os.environ['PYTHON_COVERAGE']

src/main/java/com/google/devtools/build/lib/rules/python/PyRuntime.java

@@ -37,6 +37,7 @@ public ConfiguredTarget create(RuleContext ruleContext)
 
     NestedSet<Artifact> files = PrerequisiteArtifacts.nestedSet(ruleContext, "files");
     Artifact interpreter = ruleContext.getPrerequisiteArtifact("interpreter");
+    Artifact coverageTool = ruleContext.getPrerequisiteArtifact("coverage_tool");
     PathFragment interpreterPath =
         PathFragment.create(ruleContext.attributes().get("interpreter_path", Type.STRING));
     PythonVersion pythonVersion =
@@ -81,8 +82,17 @@ public ConfiguredTarget create(RuleContext ruleContext)
 
     PyRuntimeInfo provider =
         hermetic
-            ? PyRuntimeInfo.createForInBuildRuntime(interpreter, files, pythonVersion, stubShebang)
-            : PyRuntimeInfo.createForPlatformRuntime(interpreterPath, pythonVersion, stubShebang);
+            ? PyRuntimeInfo.createForInBuildRuntime(
+                interpreter,
+                files,
+                coverageTool,
+                pythonVersion,
+                stubShebang)
+            : PyRuntimeInfo.createForPlatformRuntime(
+                interpreterPath,
+                coverageTool,
+                pythonVersion,
+                stubShebang);
 
     return new RuleConfiguredTargetBuilder(ruleContext)
         .setFilesToBuild(files)

src/main/java/com/google/devtools/build/lib/rules/python/PyRuntimeInfo.java

@@ -56,6 +56,7 @@ public final class PyRuntimeInfo implements Info, PyRuntimeInfoApi<Artifact> {
   @Nullable private final Artifact interpreter;
   // Validated on initialization to contain Artifact
   @Nullable private final Depset files;
+  @Nullable private final Artifact coverageTool;
   /** Invariant: either PY2 or PY3. */
   private final PythonVersion pythonVersion;
 
@@ -66,6 +67,7 @@ private PyRuntimeInfo(
       @Nullable PathFragment interpreterPath,
       @Nullable Artifact interpreter,
       @Nullable Depset files,
+      @Nullable Artifact coverageTool,
       PythonVersion pythonVersion,
       @Nullable String stubShebang) {
     Preconditions.checkArgument((interpreterPath == null) != (interpreter == null));
@@ -75,6 +77,7 @@ private PyRuntimeInfo(
     this.files = files;
     this.interpreterPath = interpreterPath;
     this.interpreter = interpreter;
+    this.coverageTool = coverageTool;
     this.pythonVersion = pythonVersion;
     if (stubShebang != null && !stubShebang.isEmpty()) {
       this.stubShebang = stubShebang;
@@ -97,25 +100,31 @@ public Location getCreationLocation() {
   public static PyRuntimeInfo createForInBuildRuntime(
       Artifact interpreter,
       NestedSet<Artifact> files,
+      @Nullable Artifact coverageTool,
       PythonVersion pythonVersion,
       @Nullable String stubShebang) {
     return new PyRuntimeInfo(
         /*location=*/ null,
         /*interpreterPath=*/ null,
         interpreter,
         Depset.of(Artifact.TYPE, files),
+        coverageTool,
         pythonVersion,
         stubShebang);
   }
 
   /** Constructs an instance from native rule logic (built-in location) for a platform runtime. */
   public static PyRuntimeInfo createForPlatformRuntime(
-      PathFragment interpreterPath, PythonVersion pythonVersion, @Nullable String stubShebang) {
+      PathFragment interpreterPath,
+      @Nullable Artifact coverageTool,
+      PythonVersion pythonVersion,
+      @Nullable String stubShebang) {
     return new PyRuntimeInfo(
         /*location=*/ null,
         interpreterPath,
         /*interpreter=*/ null,
         /*files=*/ null,
+        coverageTool,
         pythonVersion,
         stubShebang);
   }
@@ -131,12 +140,19 @@ public boolean equals(Object other) {
     return (this.interpreterPath.equals(otherInfo.interpreterPath)
         && this.interpreter.equals(otherInfo.interpreter)
         && this.files.equals(otherInfo.files)
+        && this.coverageTool.equals(otherInfo.coverageTool)
         && this.stubShebang.equals(otherInfo.stubShebang));
   }
 
   @Override
   public int hashCode() {
-    return Objects.hash(PyRuntimeInfo.class, interpreterPath, interpreter, files, stubShebang);
+    return Objects.hash(
+      PyRuntimeInfo.class,
+      interpreterPath,
+      interpreter,
+      coverageTool,
+      files,
+      stubShebang);
   }
 
   /**
@@ -191,6 +207,12 @@ public Depset getFilesForStarlark() {
     return files;
   }
 
+  @Override
+  @Nullable
+  public Artifact getCoverageTool() {
+    return coverageTool;
+  }
+
   public PythonVersion getPythonVersion() {
     return pythonVersion;
   }
@@ -213,6 +235,7 @@ public PyRuntimeInfo constructor(
         Object interpreterPathUncast,
         Object interpreterUncast,
         Object filesUncast,
+        Object coverageToolUncast,
         String pythonVersion,
         String stubShebang,
         StarlarkThread thread)
@@ -226,6 +249,8 @@ public PyRuntimeInfo constructor(
         Depset.cast(filesUncast, Artifact.class, "files");
         filesDepset = (Depset) filesUncast;
       }
+      Artifact coverageTool =
+          coverageToolUncast == NONE ? null : (Artifact) coverageToolUncast;
 
       if ((interpreter == null) == (interpreterPath == null)) {
         throw Starlark.errorf(
@@ -253,6 +278,7 @@ public PyRuntimeInfo constructor(
             /*interpreterPath=*/ null,
             interpreter,
             filesDepset,
+            coverageTool,
             parsedPythonVersion,
             stubShebang);
       } else {
@@ -261,6 +287,7 @@ public PyRuntimeInfo constructor(
             PathFragment.create(interpreterPath),
             /*interpreter=*/ null,
             /*files=*/ null,
+            coverageTool,
             parsedPythonVersion,
             stubShebang);
       }

src/main/java/com/google/devtools/build/lib/rules/python/PyRuntimeRule.java

@@ -55,6 +55,15 @@ public RuleClass build(RuleClass.Builder builder, RuleDefinitionEnvironment env)
         <!-- #END_BLAZE_RULE.ATTRIBUTE --> */
         .add(attr("interpreter_path", STRING))
 
+        /* <!-- #BLAZE_RULE(py_runtime).ATTRIBUTE(coverage_tool) -->
+        This is a target to use for collecting code coverage information from python tests.
+
+        <p>If set, the path to this target determine the value for the `PYTHON_COVERAGE` environment
+        variable, and will add this target to the runfiles for test runs where coverage is enabled.
+        </p>
+        <!-- #END_BLAZE_RULE.ATTRIBUTE --> */
+        .add(attr("coverage_tool", LABEL).allowedFileTypes(FileTypeSet.ANY_FILE).singleArtifact())
+
         /* <!-- #BLAZE_RULE(py_runtime).ATTRIBUTE(python_version) -->
         Whether this runtime is for Python major version 2 or 3. Valid values are <code>"PY2"</code>
         and <code>"PY3"</code>.

src/main/java/com/google/devtools/build/lib/starlarkbuildapi/python/PyRuntimeInfoApi.java

@@ -82,6 +82,16 @@ public interface PyRuntimeInfoApi<FileT extends FileApi> extends StarlarkValue {
   @Nullable
   Depset getFilesForStarlark();
 
+  @StarlarkMethod(
+      name = "coverage_tool",
+      structField = true,
+      allowReturnNones = true,
+      doc =
+          "If set, this field is a <code>File</code> representing tool used for collecting code "
+              + "coverage information from python tests. Otherwise, this is <code>None</code>.")
+  @Nullable
+  FileT getCoverageTool();
+
   @StarlarkMethod(
       name = "python_version",
       structField = true,
@@ -145,6 +155,16 @@ interface PyRuntimeInfoProviderApi extends ProviderApi {
                       + "for this argument if you pass in <code>interpreter_path</code>. If "
                       + "<code>interpreter</code> is given and this argument is <code>None</code>, "
                       + "<code>files</code> becomes an empty <code>depset</code> instead."),
+          @Param(
+              name = "coverage_tool",
+              allowedTypes = {
+                @ParamType(type = FileApi.class),
+                @ParamType(type = NoneType.class),
+              },
+              positional = false,
+              named = true,
+              defaultValue = "None",
+              doc = "The value for the new object's <code>coverage_tool</code> field."),
           @Param(
               name = "python_version",
               positional = false,
@@ -169,6 +189,7 @@ PyRuntimeInfoApi<?> constructor(
         Object interpreterPathUncast,
         Object interpreterUncast,
         Object filesUncast,
+        Object coverageToolUncast,
         String pythonVersion,
         String stubShebang,
         StarlarkThread thread)