From f75c7df9947e19deaf5bd4474f35a8bb51b0ba37 Mon Sep 17 00:00:00 2001
From: Eric Arellano <14852634+Eric-Arellano@users.noreply.github.com>
Date: Mon, 3 Jan 2022 13:16:31 -0700
Subject: [PATCH] Fix `.pyi` type stubs to show up in `python_distribution`
 (cherrypick of #14033) (#14035)

@asherf discovered that even though we were including `.pyi` files in the chroot we give to `setuptools`, we need to set the value in `package_data` for setuptools to actually include the file. See https://blog.ian.stapletoncordas.co/2019/02/distributing-python-libraries-with-type-annotations.html.

Because of this issue, `pantsbuild.pants` was not including `native_engine.pyi` in the wheel.

[ci skip-rust]
[ci skip-build-wheels]
---
 .../pants/backend/python/goals/setup_py.py    | 24 ++++++++++++-------
 .../backend/python/goals/setup_py_test.py     |  2 +-
 2 files changed, 17 insertions(+), 9 deletions(-)

diff --git a/src/python/pants/backend/python/goals/setup_py.py b/src/python/pants/backend/python/goals/setup_py.py
index a1d3548bbb9..2dc9d725ce8 100644
--- a/src/python/pants/backend/python/goals/setup_py.py
+++ b/src/python/pants/backend/python/goals/setup_py.py
@@ -847,7 +847,6 @@ def find_packages(
     """
     # Find all packages implied by all the sources.
     packages: set[str] = set()
-    package_data: DefaultDict[str, list[str]] = defaultdict(list)
     for file_path in itertools.chain(python_files, resource_files):
         # Python 2: An __init__.py file denotes a package.
         # Python 3: Any directory containing python source files is a package.
@@ -855,16 +854,25 @@ def find_packages(
             packages.add(os.path.dirname(file_path).replace(os.path.sep, "."))
 
     # Now find all package_data.
-    for resource_file in resource_files:
-        # Find the closest enclosing package, if any.  Resources will be loaded relative to that.
-        maybe_package: str = os.path.dirname(resource_file).replace(os.path.sep, ".")
+    package_data: DefaultDict[str, list[str]] = defaultdict(list)
+
+    def maybe_add_resource(fp: str) -> None:
+        # Find the closest enclosing package, if any. Resources will be loaded relative to that.
+        maybe_package: str = os.path.dirname(fp).replace(os.path.sep, ".")
         while maybe_package and maybe_package not in packages:
             maybe_package = maybe_package.rpartition(".")[0]
         # If resource is not in a package, ignore it. There's no principled way to load it anyway.
-        if maybe_package:
-            package_data[maybe_package].append(
-                os.path.relpath(resource_file, maybe_package.replace(".", os.path.sep))
-            )
+        if not maybe_package:
+            return
+        package_data[maybe_package].append(
+            os.path.relpath(fp, maybe_package.replace(".", os.path.sep))
+        )
+
+    for resource_file in resource_files:
+        maybe_add_resource(resource_file)
+    for py_file in python_files:
+        if py_file.endswith(".pyi"):
+            maybe_add_resource(py_file)
 
     # See which packages are pkg_resources-style namespace packages.
     # Note that implicit PEP 420 namespace packages and pkgutil-style namespace packages
diff --git a/src/python/pants/backend/python/goals/setup_py_test.py b/src/python/pants/backend/python/goals/setup_py_test.py
index 78ba33bbed2..28f43eb720c 100644
--- a/src/python/pants/backend/python/goals/setup_py_test.py
+++ b/src/python/pants/backend/python/goals/setup_py_test.py
@@ -407,7 +407,7 @@ def test_generate_chroot(chroot_rule_runner: RuleRunner) -> None:
             "plugin_demo": "hello world",
             "packages": ("foo", "foo.qux"),
             "namespace_packages": ("foo",),
-            "package_data": {"foo": ("resources/js/code.js",)},
+            "package_data": {"foo": ("resources/js/code.js",), "foo.qux": ("qux.pyi",)},
             "install_requires": ("baz==1.1.1",),
             "entry_points": {"console_scripts": ["foo_main = foo.qux.bin:main"]},
         },