From 1fb161233ed41d5318f993c38a03a284fd1cb028 Mon Sep 17 00:00:00 2001
From: Matt Mackay <matt@aspect.dev>
Date: Mon, 31 Jul 2023 17:25:18 -0400
Subject: [PATCH] fix: ensure coverage providers are output from both
 py_library and py_binary rules (#168)

---
 docs/py_library.md        | 37 +++++++++++++++++++++++++++++--------
 py/private/py_binary.bzl  |  5 +++++
 py/private/py_library.bzl | 18 +++++++++++++++---
 3 files changed, 49 insertions(+), 11 deletions(-)

diff --git a/docs/py_library.md b/docs/py_library.md
index 2b9c1b7f..9712656d 100644
--- a/docs/py_library.md
+++ b/docs/py_library.md
@@ -24,12 +24,12 @@ py_library(<a href="#py_library-name">name</a>, <a href="#py_library-data">data<
 | <a id="py_library-srcs"></a>srcs |  -   | <a href="https://bazel.build/concepts/labels">List of labels</a> | optional | <code>[]</code> |
 
 
-<a id="py_library_utils.make_srcs_depset"></a>
+<a id="py_library_utils.implementation"></a>
 
-## py_library_utils.make_srcs_depset
+## py_library_utils.implementation
 
 <pre>
-py_library_utils.make_srcs_depset(<a href="#py_library_utils.make_srcs_depset-ctx">ctx</a>)
+py_library_utils.implementation(<a href="#py_library_utils.implementation-ctx">ctx</a>)
 </pre>
 
 
@@ -39,7 +39,7 @@ py_library_utils.make_srcs_depset(<a href="#py_library_utils.make_srcs_depset-ct
 
 | Name  | Description | Default Value |
 | :------------- | :------------- | :------------- |
-| <a id="py_library_utils.make_srcs_depset-ctx"></a>ctx |  <p align="center"> - </p>   |  none |
+| <a id="py_library_utils.implementation-ctx"></a>ctx |  <p align="center"> - </p>   |  none |
 
 
 <a id="py_library_utils.make_imports_depset"></a>
@@ -60,6 +60,27 @@ py_library_utils.make_imports_depset(<a href="#py_library_utils.make_imports_dep
 | <a id="py_library_utils.make_imports_depset-ctx"></a>ctx |  <p align="center"> - </p>   |  none |
 
 
+<a id="py_library_utils.make_instrumented_files_info"></a>
+
+## py_library_utils.make_instrumented_files_info
+
+<pre>
+py_library_utils.make_instrumented_files_info(<a href="#py_library_utils.make_instrumented_files_info-ctx">ctx</a>, <a href="#py_library_utils.make_instrumented_files_info-extra_source_attributes">extra_source_attributes</a>,
+                                              <a href="#py_library_utils.make_instrumented_files_info-extra_dependency_attributes">extra_dependency_attributes</a>)
+</pre>
+
+
+
+**PARAMETERS**
+
+
+| Name  | Description | Default Value |
+| :------------- | :------------- | :------------- |
+| <a id="py_library_utils.make_instrumented_files_info-ctx"></a>ctx |  <p align="center"> - </p>   |  none |
+| <a id="py_library_utils.make_instrumented_files_info-extra_source_attributes"></a>extra_source_attributes |  <p align="center"> - </p>   |  <code>[]</code> |
+| <a id="py_library_utils.make_instrumented_files_info-extra_dependency_attributes"></a>extra_dependency_attributes |  <p align="center"> - </p>   |  <code>[]</code> |
+
+
 <a id="py_library_utils.make_merged_runfiles"></a>
 
 ## py_library_utils.make_merged_runfiles
@@ -81,12 +102,12 @@ py_library_utils.make_merged_runfiles(<a href="#py_library_utils.make_merged_run
 | <a id="py_library_utils.make_merged_runfiles-extra_runfiles_depsets"></a>extra_runfiles_depsets |  <p align="center"> - </p>   |  <code>[]</code> |
 
 
-<a id="py_library_utils.implementation"></a>
+<a id="py_library_utils.make_srcs_depset"></a>
 
-## py_library_utils.implementation
+## py_library_utils.make_srcs_depset
 
 <pre>
-py_library_utils.implementation(<a href="#py_library_utils.implementation-ctx">ctx</a>)
+py_library_utils.make_srcs_depset(<a href="#py_library_utils.make_srcs_depset-ctx">ctx</a>)
 </pre>
 
 
@@ -96,6 +117,6 @@ py_library_utils.implementation(<a href="#py_library_utils.implementation-ctx">c
 
 | Name  | Description | Default Value |
 | :------------- | :------------- | :------------- |
-| <a id="py_library_utils.implementation-ctx"></a>ctx |  <p align="center"> - </p>   |  none |
+| <a id="py_library_utils.make_srcs_depset-ctx"></a>ctx |  <p align="center"> - </p>   |  none |
 
 
diff --git a/py/private/py_binary.bzl b/py/private/py_binary.bzl
index f5ac44de..cd435aba 100644
--- a/py/private/py_binary.bzl
+++ b/py/private/py_binary.bzl
@@ -74,6 +74,10 @@ def _py_binary_rule_imp(ctx):
     )
 
     imports = _py_library.make_imports_depset(ctx)
+    instrumented_files_info = _py_library.make_instrumented_files_info(
+        ctx, 
+        extra_source_attributes = ["main"]
+    )
 
     return [
         DefaultInfo(
@@ -88,6 +92,7 @@ def _py_binary_rule_imp(ctx):
             has_py3_only_sources = True,
             uses_shared_libraries = False,
         ),
+        instrumented_files_info,
     ]
 
 _attrs = dict({
diff --git a/py/private/py_library.bzl b/py/private/py_library.bzl
index 332569ac..370a53ac 100644
--- a/py/private/py_library.bzl
+++ b/py/private/py_library.bzl
@@ -4,6 +4,14 @@ load("@bazel_skylib//lib:paths.bzl", "paths")
 load("//py/private:providers.bzl", "PyWheelInfo")
 load("//py/private:py_wheel.bzl", py_wheel = "py_wheel_lib")
 
+def _make_instrumented_files_info(ctx, extra_source_attributes = [], extra_dependency_attributes = []):
+    return coverage_common.instrumented_files_info(
+        ctx,
+        source_attributes = ["srcs"] + extra_source_attributes,
+        dependency_attributes = ["data", "deps"] + extra_dependency_attributes,
+        extensions = ["py"],
+    )
+
 def _make_srcs_depset(ctx):
     return depset(
         order = "postorder",
@@ -89,6 +97,7 @@ def _py_library_impl(ctx):
     transitive_srcs = _make_srcs_depset(ctx)
     imports = _make_imports_depset(ctx)
     runfiles = _make_merged_runfiles(ctx)
+    instrumented_files_info = _make_instrumented_files_info(ctx)
     py_wheel_info = py_wheel.make_py_wheel_info(ctx, ctx.attr.deps)
 
     return [
@@ -104,6 +113,7 @@ def _py_library_impl(ctx):
             uses_shared_libraries = False,
         ),
         py_wheel_info,
+        instrumented_files_info,
     ]
 
 _attrs = dict({
@@ -126,11 +136,13 @@ _providers = [
 ]
 
 py_library_utils = struct(
-    make_srcs_depset = _make_srcs_depset,
+    # keep-sorted
+    attrs = _attrs,
+    implementation = _py_library_impl,
     make_imports_depset = _make_imports_depset,
+    make_instrumented_files_info = _make_instrumented_files_info,
     make_merged_runfiles = _make_merged_runfiles,
-    implementation = _py_library_impl,
-    attrs = _attrs,
+    make_srcs_depset = _make_srcs_depset,
     py_library_providers = _providers,
 )