feat: Package pyi files in wheel (#2609)

1.1.0 introduced separate attributes for the type definitions (`.pyi`
files) and type checking. This patch adds those files to the wheel to
ensure that they are distributed and available to users.

#2538 introduced
`pyi_srcs`.

---------

Co-authored-by: Ignas Anikevicius <[email protected]>

CHANGELOG.md

@@ -53,6 +53,8 @@ Unreleased changes template.
 {#v0-0-0-changed}
 ### Changed
 * (deps) platforms 0.0.4 -> 0.0.11
+* (py_wheel) Package `py_library.pyi_srcs` (`.pyi` files) in the wheel.
+* (py_package) Package `py_library.pyi_srcs` (`.pyi` files) in `py_package`.
 
 {#v0-0-0-fixed}
 ### Fixed

examples/wheel/BUILD.bazel

@@ -33,6 +33,7 @@ py_library(
     deps = [
         "//examples/wheel/lib:simple_module",
         "//examples/wheel/lib:module_with_data",
+        "//examples/wheel/lib:module_with_type_annotations",
         # Example dependency which is not packaged in the wheel
         # due to "packages" filter on py_package rule.
         "//tests/load_from_macro:foo",
@@ -67,6 +68,7 @@ py_wheel(
     version = "0.0.1",
     deps = [
         "//examples/wheel/lib:module_with_data",
+        "//examples/wheel/lib:module_with_type_annotations",
         "//examples/wheel/lib:simple_module",
     ],
 )
@@ -90,6 +92,7 @@ py_wheel(
     version = "$(VERSION)",
     deps = [
         "//examples/wheel/lib:module_with_data",
+        "//examples/wheel/lib:module_with_type_annotations",
         "//examples/wheel/lib:simple_module",
     ],
 )
@@ -109,6 +112,7 @@ py_wheel(
     version = "0.1.{BUILD_TIMESTAMP}",
     deps = [
         "//examples/wheel/lib:module_with_data",
+        "//examples/wheel/lib:module_with_type_annotations",
         "//examples/wheel/lib:simple_module",
     ],
 )

examples/wheel/lib/BUILD.bazel

@@ -23,6 +23,12 @@ py_library(
     srcs = ["simple_module.py"],
 )
 
+py_library(
+    name = "module_with_type_annotations",
+    srcs = ["module_with_type_annotations.py"],
+    pyi_srcs = ["module_with_type_annotations.pyi"],
+)
+
 py_library(
     name = "module_with_data",
     srcs = ["module_with_data.py"],

examples/wheel/lib/module_with_type_annotations.py

@@ -0,0 +1,16 @@
+# Copyright 2025 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.
+
+def function():
+    return "qux"

examples/wheel/lib/module_with_type_annotations.pyi

@@ -0,0 +1,15 @@
+# Copyright 2025 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.
+
+def function() -> str: ...

examples/wheel/main.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import examples.wheel.lib.module_with_data as module_with_data
+import examples.wheel.lib.module_with_type_annotations as module_with_type_annotations
 import examples.wheel.lib.simple_module as simple_module
 
 
@@ -23,6 +24,7 @@ def function():
 def main():
     print(function())
     print(module_with_data.function())
+    print(module_with_type_annotations.function())
     print(simple_module.function())
 
 

examples/wheel/test_publish.py

@@ -104,7 +104,7 @@ def test_upload_and_query_simple_api(self):
     </head>
     <body>
         <h1>Links for example-minimal-library</h1>
-             <a href="/packages/example_minimal_library-0.0.1-py3-none-any.whl#sha256=79a4e9c1838c0631d5d8fa49a26efd6e9a364f6b38d9597c0f6df112271a0e28">example_minimal_library-0.0.1-py3-none-any.whl</a><br>
+             <a href="/packages/example_minimal_library-0.0.1-py3-none-any.whl#sha256=0cbf4ec574676015af595f570caf4ae2812f994f6338e247b002b4e496b6fbd5">example_minimal_library-0.0.1-py3-none-any.whl</a><br>
     </body>
 </html>"""
             self.assertEqual(

examples/wheel/wheel_test.py

@@ -76,14 +76,16 @@ def test_py_library_wheel(self):
                 zf.namelist(),
                 [
                     "examples/wheel/lib/module_with_data.py",
+                    "examples/wheel/lib/module_with_type_annotations.py",
+                    "examples/wheel/lib/module_with_type_annotations.pyi",
                     "examples/wheel/lib/simple_module.py",
                     "example_minimal_library-0.0.1.dist-info/WHEEL",
                     "example_minimal_library-0.0.1.dist-info/METADATA",
                     "example_minimal_library-0.0.1.dist-info/RECORD",
                 ],
             )
         self.assertFileSha256Equal(
-            filename, "79a4e9c1838c0631d5d8fa49a26efd6e9a364f6b38d9597c0f6df112271a0e28"
+            filename, "0cbf4ec574676015af595f570caf4ae2812f994f6338e247b002b4e496b6fbd5"
         )
 
     def test_py_package_wheel(self):
@@ -98,6 +100,8 @@ def test_py_package_wheel(self):
                     "examples/wheel/lib/data,with,commas.txt",
                     "examples/wheel/lib/data.txt",
                     "examples/wheel/lib/module_with_data.py",
+                    "examples/wheel/lib/module_with_type_annotations.py",
+                    "examples/wheel/lib/module_with_type_annotations.pyi",
                     "examples/wheel/lib/simple_module.py",
                     "examples/wheel/main.py",
                     "example_minimal_package-0.0.1.dist-info/WHEEL",
@@ -106,7 +110,7 @@ def test_py_package_wheel(self):
                 ],
             )
         self.assertFileSha256Equal(
-            filename, "82370bf61310e2d3c7b1218368457dc7e161bf5dc1a280d7d45102b5e56acf43"
+            filename, "22aff90dd3c8c30c3ce2b729bb793cab0bd2668a6810de232677a0354ce79cae"
         )
 
     def test_customized_wheel(self):
@@ -121,6 +125,8 @@ def test_customized_wheel(self):
                     "examples/wheel/lib/data,with,commas.txt",
                     "examples/wheel/lib/data.txt",
                     "examples/wheel/lib/module_with_data.py",
+                    "examples/wheel/lib/module_with_type_annotations.py",
+                    "examples/wheel/lib/module_with_type_annotations.pyi",
                     "examples/wheel/lib/simple_module.py",
                     "examples/wheel/main.py",
                     "example_customized-0.0.1.dist-info/WHEEL",
@@ -145,8 +151,10 @@ def test_customized_wheel(self):
 "examples/wheel/lib/data,with,commas.txt",sha256=9vJKEdfLu8bZRArKLroPZJh1XKkK3qFMXiM79MBL2Sg,12
 examples/wheel/lib/data.txt,sha256=9vJKEdfLu8bZRArKLroPZJh1XKkK3qFMXiM79MBL2Sg,12
 examples/wheel/lib/module_with_data.py,sha256=8s0Khhcqz3yVsBKv2IB5u4l4TMKh7-c_V6p65WVHPms,637
+examples/wheel/lib/module_with_type_annotations.py,sha256=MM2cFQsCBaUnzGiEGT5r07jhKSaCVRh5Paw_YLyrS-w,636
+examples/wheel/lib/module_with_type_annotations.pyi,sha256=fja3ql_WRJ1qO8jyZjWWrTTMcg1J7EpOQivOHY_8vI4,630
 examples/wheel/lib/simple_module.py,sha256=z2hwciab_XPNIBNH8B1Q5fYgnJvQTeYf0ZQJpY8yLLY,637
-examples/wheel/main.py,sha256=sgg5iWN_9inYBjm6_Zw27hYdmo-l24fA-2rfphT-IlY,909
+examples/wheel/main.py,sha256=mFiRfzQEDwCHr-WVNQhOH26M42bw1UMF6IoqvtuDTrw,1047
 example_customized-0.0.1.dist-info/WHEEL,sha256=sobxWSyDDkdg_rinUth-jxhXHqoNqlmNMJY3aTZn2Us,91
 example_customized-0.0.1.dist-info/METADATA,sha256=QYQcDJFQSIqan8eiXqL67bqsUfgEAwf2hoK_Lgi1S-0,559
 example_customized-0.0.1.dist-info/entry_points.txt,sha256=pqzpbQ8MMorrJ3Jp0ntmpZcuvfByyqzMXXi2UujuXD0,137
@@ -197,7 +205,7 @@ def test_customized_wheel(self):
 second = second.main:s""",
             )
         self.assertFileSha256Equal(
-            filename, "706e8dd45884d8cb26e92869f7d29ab7ed9f683b4e2d08f06c03dbdaa12191b8"
+            filename, "657a938a6fdd6f38bf73d1d91016ffff85d68cf29ca390692a3e9d923dd0e39e"
         )
 
     def test_filename_escaping(self):
@@ -211,6 +219,8 @@ def test_filename_escaping(self):
                     "examples/wheel/lib/data,with,commas.txt",
                     "examples/wheel/lib/data.txt",
                     "examples/wheel/lib/module_with_data.py",
+                    "examples/wheel/lib/module_with_type_annotations.py",
+                    "examples/wheel/lib/module_with_type_annotations.pyi",
                     "examples/wheel/lib/simple_module.py",
                     "examples/wheel/main.py",
                     # PEP calls for replacing only in the archive filename.
@@ -248,6 +258,8 @@ def test_custom_package_root_wheel(self):
                     "wheel/lib/data,with,commas.txt",
                     "wheel/lib/data.txt",
                     "wheel/lib/module_with_data.py",
+                    "wheel/lib/module_with_type_annotations.py",
+                    "wheel/lib/module_with_type_annotations.pyi",
                     "wheel/lib/simple_module.py",
                     "wheel/main.py",
                     "examples_custom_package_root-0.0.1.dist-info/WHEEL",
@@ -265,7 +277,7 @@ def test_custom_package_root_wheel(self):
             for line in record_contents.splitlines():
                 self.assertFalse(line.startswith("/"))
         self.assertFileSha256Equal(
-            filename, "568922541703f6edf4b090a8413991f9fa625df2844e644dd30bdbe9deb660be"
+            filename, "d415edbf8f326161674c1fa260e364dd44f2a0311e2f596284320ea52d2a8bdb"
         )
 
     def test_custom_package_root_multi_prefix_wheel(self):
@@ -281,6 +293,8 @@ def test_custom_package_root_multi_prefix_wheel(self):
                     "data,with,commas.txt",
                     "data.txt",
                     "module_with_data.py",
+                    "module_with_type_annotations.py",
+                    "module_with_type_annotations.pyi",
                     "simple_module.py",
                     "main.py",
                     "example_custom_package_root_multi_prefix-0.0.1.dist-info/WHEEL",
@@ -297,7 +311,7 @@ def test_custom_package_root_multi_prefix_wheel(self):
             for line in record_contents.splitlines():
                 self.assertFalse(line.startswith("/"))
         self.assertFileSha256Equal(
-            filename, "a8b91ce9d6f570e97b40a357a292a6f595d3470f07c479cb08550257cc9c8306"
+            filename, "6b76a1178c90996feaf3f9417f350c4a67f90f4247647fd4fd552858dc372d4b"
         )
 
     def test_custom_package_root_multi_prefix_reverse_order_wheel(self):
@@ -313,6 +327,8 @@ def test_custom_package_root_multi_prefix_reverse_order_wheel(self):
                     "lib/data,with,commas.txt",
                     "lib/data.txt",
                     "lib/module_with_data.py",
+                    "lib/module_with_type_annotations.py",
+                    "lib/module_with_type_annotations.pyi",
                     "lib/simple_module.py",
                     "main.py",
                     "example_custom_package_root_multi_prefix_reverse_order-0.0.1.dist-info/WHEEL",
@@ -329,7 +345,7 @@ def test_custom_package_root_multi_prefix_reverse_order_wheel(self):
             for line in record_contents.splitlines():
                 self.assertFalse(line.startswith("/"))
         self.assertFileSha256Equal(
-            filename, "8f44e940731757c186079a42cfe7ea3d43cd96b526e3fb2ca2a3ea3048a9d489"
+            filename, "f976f0bb1c7d753e8c41629d6b79fb09908c6ecd2fec006816879fc86b664f3f"
         )
 
     def test_python_requires_wheel(self):
@@ -354,7 +370,7 @@ def test_python_requires_wheel(self):
 """,
             )
         self.assertFileSha256Equal(
-            filename, "ba32493f5e43e481346384aaab9e8fa09c23884276ad057c5f432096a0350101"
+            filename, "f3b74ce429c3324b87f8d1cc7dc33be1493f54bb88d546a7d53be7587b82c1a7"
         )
 
     def test_python_abi3_binary_wheel(self):
@@ -419,7 +435,7 @@ def test_rule_creates_directory_and_is_included_in_wheel(self):
                 ],
             )
         self.assertFileSha256Equal(
-            filename, "ac9216bd54dcae1a6270c35fccf8a73b0be87c1b026c28e963b7c76b2f9b722b"
+            filename, "d8e874b807e5574bd11a9312c58ce7fe7055afb80412d0d0e7ed21fc9223cd53"
         )
 
     def test_rule_expands_workspace_status_keys_in_wheel_metadata(self):

python/private/py_package.bzl

@@ -46,6 +46,9 @@ def _py_package_impl(ctx):
     if hasattr(py_info, "transitive_pyc_files"):
         inputs.add(py_info.transitive_pyc_files)
 
+    if hasattr(py_info, "transitive_pyi_files"):
+        inputs.add(py_info.transitive_pyi_files)
+
     inputs = inputs.build()
 
     # TODO: '/' is wrong on windows, but the path separator is not available in starlark.

python/private/py_wheel.bzl

@@ -14,6 +14,7 @@
 
 "Implementation of py_wheel rule"
 
+load(":py_info.bzl", "PyInfo")
 load(":py_package.bzl", "py_package_lib")
 load(":py_wheel_normalize_pep440.bzl", "normalize_pep440")
 load(":stamp.bzl", "is_stamping_enabled")
@@ -319,8 +320,13 @@ def _py_wheel_impl(ctx):
 
     name_file = ctx.actions.declare_file(ctx.label.name + ".name")
 
+    direct_pyi_files = []
+    for dep in ctx.attr.deps:
+        if PyInfo in dep:
+            direct_pyi_files.extend(dep[PyInfo].direct_pyi_files.to_list())
+
     inputs_to_package = depset(
-        direct = ctx.files.deps,
+        direct = ctx.files.deps + direct_pyi_files,
     )
 
     # Inputs to this rule which are not to be packaged.

tests/whl_filegroup/extract_wheel_files_test.py

@@ -14,6 +14,8 @@ def test_get_wheel_record(self) -> None:
             "examples/wheel/lib/data,with,commas.txt",
             "examples/wheel/lib/data.txt",
             "examples/wheel/lib/module_with_data.py",
+            "examples/wheel/lib/module_with_type_annotations.py",
+            "examples/wheel/lib/module_with_type_annotations.pyi",
             "examples/wheel/lib/simple_module.py",
             "examples/wheel/main.py",
             "example_minimal_package-0.0.1.dist-info/WHEEL",