Generate .pyi files in py_proto_library (#10366) (#21567)

We here adjust py_proto_library to produce `.pyi` files along with the `.py` files it already generates. This achieves the same effect as the gRPC py_proto_library. Per reviewer feedback, we here propagate the paths to the `.pyi` files via the `direct_pyi_files` and `transitive_pyi_files` members of `PyInfo` which are available starting in `rules_python` 1.1.0. Since this means that the files are not automatically in the runfiles of the `py_proto_library`, tooling to run mypy will need to be adjusted to add what is needed from these provider fields to the runfiles if that is what is wanted. Closes#10366. Closes#21567 COPYBARA_INTEGRATE_REVIEW=#21567 from dws:py_proto_library-pyi 37fbdc0 PiperOrigin-RevId: 778116905

Author: dws

bazel/py_proto_library.bzl

@@ -17,6 +17,13 @@ _PyProtoInfo = provider(
  (depset[File]) Files from the transitive closure implicit proto
  dependencies""",
 "transitive_sources": """(depset[File]) The Python sources.""",
+"direct_pyi_files": """
+ (depset[File]) Type definition files (usually `.pyi` files)
+ for the Python modules provided by this target.""",
+"transitive_pyi_files": """
+ (depset[File]) The transitive set of type definition files
+ (usually `.pyi` files) for the Python modules for this target
+ and its transitive dependencies.""",
  },
 )
 
@@ -61,6 +68,7 @@ def _py_proto_aspect_impl(target, ctx):
 api_deps= [proto_lang_toolchain_info.runtime]
 
 generated_sources= []
+generated_stubs= []
 proto_info=target[ProtoInfo]
 proto_root=proto_info.proto_source_root
 ifproto_info.direct_sources:
@@ -72,6 +80,14 @@ def _py_proto_aspect_impl(target, ctx):
 name_mapper=lambdaname: name.replace("-", "_").replace(".", "/"),
  )
 
+# Generate pyi files
+generated_stubs=proto_common.declare_generated_files(
+actions=ctx.actions,
+proto_info=proto_info,
+extension="_pb2.pyi",
+name_mapper=lambdaname: name.replace("-", "_").replace(".", "/"),
+ )
+
 # Handles multiple repository and virtual import cases
 ifproto_root.startswith(ctx.bin_dir.path):
 proto_root=proto_root[len(ctx.bin_dir.path) +1:]
@@ -84,12 +100,17 @@ def _py_proto_aspect_impl(target, ctx):
 else:
 proto_root=ctx.workspace_name+"/"+proto_root
 
+additional_args=ctx.actions.args()
+ifgenerated_stubs:
+additional_args.add(plugin_output, format="--pyi_out=%s")
+
 proto_common.compile(
 actions=ctx.actions,
 proto_info=proto_info,
 proto_lang_toolchain_info=proto_lang_toolchain_info,
-generated_files=generated_sources,
+generated_files=generated_sources+generated_stubs,
 plugin_output=plugin_output,
+additional_args=additional_args,
  )
 
 # Generated sources == Python sources
@@ -104,6 +125,13 @@ def _py_proto_aspect_impl(target, ctx):
 direct=python_sources,
 transitive= [dep.transitive_sourcesfordepindeps],
  )
+direct_pyi_files=depset(
+direct=generated_stubs,
+ )
+transitive_pyi_files=depset(
+direct=generated_stubs,
+transitive= [dep.transitive_pyi_filesfordepindeps],
+ )
 
 return [
 _PyProtoInfo(
@@ -119,6 +147,8 @@ def _py_proto_aspect_impl(target, ctx):
  ),
 runfiles_from_proto_deps=runfiles_from_proto_deps,
 transitive_sources=transitive_sources,
+direct_pyi_files=direct_pyi_files,
+transitive_pyi_files=transitive_pyi_files,
  ),
  ]
 
@@ -150,6 +180,13 @@ def _py_proto_library_rule(ctx):
 default_outputs=depset(
 transitive= [info.transitive_sourcesforinfoinpyproto_infos],
  )
+direct_pyi_files= []
+forinfoinpyproto_infos:
+direct_pyi_files.extend(info.direct_pyi_files.to_list())
+transitive_pyi_files=depset(
+direct=direct_pyi_files,
+transitive= [info.transitive_pyi_filesforinfoinpyproto_infos],
+ )
 
 return [
 DefaultInfo(
@@ -166,6 +203,8 @@ def _py_proto_library_rule(ctx):
 PyInfo(
 transitive_sources=default_outputs,
 imports=depset(transitive= [info.importsforinfoinpyproto_infos]),
+direct_pyi_files=depset(direct=direct_pyi_files),
+transitive_pyi_files=transitive_pyi_files,
 # Proto always produces 2- and 3- compatible source files
 has_py2_only_sources=False,
 has_py3_only_sources=False,