Refactor collect_package_data to make logic clearer

stub_uploader/build_wheel.py

@@ -23,11 +23,11 @@
 import argparse
 import os
 import os.path
-from pathlib import Path
 import shutil
 import subprocess
 import sys
 import tempfile
+from pathlib import Path
 from textwrap import dedent
 from typing import Optional
 
@@ -128,6 +128,30 @@ def __init__(self, typeshed_dir: str, distribution: str) -> None:
         self.stub_dir = Path(typeshed_dir) / THIRD_PARTY_NAMESPACE / distribution
 
 
+class PackageData:
+    """Information about the packages of a distribution and their contents."""
+
+    def __init__(self, base_path: Path, package_data: dict[str, list[str]]) -> None:
+        self.base_path = base_path
+        self.package_data = package_data
+
+    @property
+    def top_level_packages(self) -> list[str]:
+        """Top level package names.
+
+        These are the packages that are not subpackages of any other package
+        and includes namespace packages.
+        """
+        return list(self.package_data.keys())
+
+    def add_file(self, package: str, filename: str, file_contents: str) -> None:
+        """Add a file to a package."""
+        entry_path = self.base_path / package
+        entry_path.mkdir(exist_ok=True)
+        (entry_path / filename).write_text(file_contents)
+        self.package_data[package].append(filename)
+
+
 def find_stub_files(top: str) -> list[str]:
     """Find all stub files for a given package, relative to package root.
 
@@ -197,49 +221,50 @@ def copy_changelog(distribution: str, dst: str) -> None:
         pass  # Ignore missing changelogs
 
 
-def collect_setup_entries(base_dir: str) -> dict[str, list[str]]:
+def collect_package_data(base_path: Path) -> PackageData:
     """Generate package data for a setuptools.setup() call.
 
     This reflects the transformations done during copying in copy_stubs().
     """
     package_data: dict[str, list[str]] = {}
-    for entry in os.listdir(base_dir):
-        if entry == META:
-            # Metadata file entry is added at the end.
+    for entry in base_path.iterdir():
+        if is_ignored_distribution_file(entry):
             continue
-        original_entry = entry
-        if os.path.isfile(os.path.join(base_dir, entry)):
-            if not entry.endswith(".pyi"):
-                if not entry.endswith((".md", ".rst")):
-                    if (
-                        subprocess.run(
-                            ["git", "check-ignore", entry], cwd=base_dir
-                        ).returncode
-                        != 0
-                    ):
-                        raise ValueError(f"Only stub files are allowed, not {entry!r}")
-                continue
-            entry = entry.split(".")[0] + SUFFIX
+        elif entry.is_file() and entry.suffix == ".pyi":
+            pkg_name = entry.stem
             # Module -> package transformation is done while copying.
-            package_data[entry] = ["__init__.pyi"]
+            stub_files = ["__init__.pyi"]
+        elif entry.is_dir():
+            pkg_name = entry.name
+            stub_files = find_stub_files(str(entry))
         else:
-            if entry == TESTS_NAMESPACE:
-                continue
-            entry += SUFFIX
-            package_data[entry] = find_stub_files(
-                os.path.join(base_dir, original_entry)
-            )
-        package_data[entry].append(META)
-    return package_data
-
-
-def add_partial_marker(package_data: dict[str, list[str]], stub_dir: str) -> None:
-    for entry, files in package_data.items():
-        entry_path = os.path.join(stub_dir, entry)
-        os.makedirs(entry_path, exist_ok=True)
-        with open(os.path.join(entry_path, "py.typed"), "w") as py_typed:
-            py_typed.write("partial\n")
-        files.append("py.typed")
+            raise ValueError(f"Only stub files are allowed, not {entry.name!r}")
+        package_data[pkg_name + SUFFIX] = [*stub_files, META]
+    return PackageData(base_path, package_data)
+
+
+def is_ignored_distribution_file(entry: Path) -> bool:
+    if entry.is_file():
+        if entry.name == META:
+            return True
+        if entry.suffix in (".md", ".rst"):
+            return True
+        if (
+            subprocess.run(
+                ["git", "check-ignore", entry.name], cwd=str(entry.parent)
+            ).returncode
+            == 0
+        ):
+            return True
+    elif entry.is_dir():
+        if entry.name == TESTS_NAMESPACE:
+            return True
+    return False
+
+
+def add_partial_markers(pkg_data: PackageData) -> None:
+    for package in pkg_data.top_level_packages:
+        pkg_data.add_file(package, "py.typed", "partial\n")
 
 
 def generate_setup_file(
@@ -253,9 +278,9 @@ def generate_setup_file(
     all_requirements = [
         str(req) for req in metadata.requires_typeshed + metadata.requires_external
     ]
-    package_data = collect_setup_entries(str(build_data.stub_dir))
+    pkg_data = collect_package_data(build_data.stub_dir)
     if metadata.partial:
-        add_partial_marker(package_data, str(build_data.stub_dir))
+        add_partial_markers(pkg_data)
     requires_python = (
         metadata.requires_python
         if metadata.requires_python is not None
@@ -269,8 +294,8 @@ def generate_setup_file(
         ),
         version=version,
         requires=all_requirements,
-        packages=list(package_data.keys()),
-        package_data=package_data,
+        packages=pkg_data.top_level_packages,
+        package_data=pkg_data.package_data,
         requires_python=requires_python,
     )
 

tests/test_unit.py

@@ -1,20 +1,18 @@
 """Unit tests for simple helpers should go here."""
 
 import datetime
-from io import StringIO
 import os
 import tempfile
+from io import StringIO
+from pathlib import Path
 from typing import Any
 
 import pytest
 from packaging.version import Version
 
-from stub_uploader.build_wheel import collect_setup_entries
-from stub_uploader.get_version import (
-    compute_stub_version,
-    ensure_specificity,
-)
-from stub_uploader.metadata import _UploadedPackages, strip_types_prefix, Metadata
+from stub_uploader.build_wheel import collect_package_data
+from stub_uploader.get_version import compute_stub_version, ensure_specificity
+from stub_uploader.metadata import Metadata, _UploadedPackages, strip_types_prefix
 from stub_uploader.ts_data import parse_requirements
 
 
@@ -99,13 +97,17 @@ def test_compute_stub_version() -> None:
     )
 
 
-def test_collect_setup_entries() -> None:
-    stubs = os.path.join("data", "test_typeshed", "stubs")
-    entries = collect_setup_entries(os.path.join(stubs, "singlefilepkg"))
-    assert entries == ({"singlefilepkg-stubs": ["__init__.pyi", "METADATA.toml"]})
+def test_collect_package_data() -> None:
+    stubs = Path("data") / "test_typeshed" / "stubs"
+    pkg_data = collect_package_data(stubs / "singlefilepkg")
+    assert pkg_data.top_level_packages == ["singlefilepkg-stubs"]
+    assert pkg_data.package_data == (
+        {"singlefilepkg-stubs": ["__init__.pyi", "METADATA.toml"]}
+    )
 
-    entries = collect_setup_entries(os.path.join(stubs, "multifilepkg"))
-    assert entries == (
+    pkg_data = collect_package_data(stubs / "multifilepkg")
+    assert pkg_data.top_level_packages == ["multifilepkg-stubs"]
+    assert pkg_data.package_data == (
         {
             "multifilepkg-stubs": [
                 "__init__.pyi",
@@ -119,8 +121,9 @@ def test_collect_setup_entries() -> None:
         }
     )
 
-    entries = collect_setup_entries(os.path.join(stubs, "nspkg"))
-    assert entries == (
+    pkg_data = collect_package_data(stubs / "nspkg")
+    assert pkg_data.top_level_packages == ["nspkg-stubs"]
+    assert pkg_data.package_data == (
         {
             "nspkg-stubs": [
                 os.path.join("innerpkg", "__init__.pyi"),
@@ -130,25 +133,25 @@ def test_collect_setup_entries() -> None:
     )
 
 
-def test_collect_setup_entries_bogusfile() -> None:
-    stubs = os.path.join("data", "test_typeshed", "stubs")
+def test_collect_package_data_bogusfile() -> None:
+    stubs = Path("data") / "test_typeshed" / "stubs"
     with pytest.raises(
         ValueError, match="Only stub files are allowed, not 'bogusfile.txt'"
     ):
-        collect_setup_entries(os.path.join(stubs, "bogusfiles"))
+        collect_package_data(stubs / "bogusfiles")
 
     # Make sure gitignored files aren't collected, nor do they crash function
     with open(os.path.join(stubs, "singlefilepkg", ".METADATA.toml.swp"), "w"):
         pass
-    entries = collect_setup_entries(os.path.join(stubs, "singlefilepkg"))
-    assert len(entries["singlefilepkg-stubs"]) == 2
+    pkg_data = collect_package_data(stubs / "singlefilepkg")
+    assert len(pkg_data.package_data["singlefilepkg-stubs"]) == 2
 
     with open(
         os.path.join(stubs, "multifilepkg", "multifilepkg", ".METADATA.toml.swp"), "w"
     ):
         pass
-    entries = collect_setup_entries(os.path.join(stubs, "multifilepkg"))
-    assert len(entries["multifilepkg-stubs"]) == 7
+    pkg_data = collect_package_data(stubs / "multifilepkg")
+    assert len(pkg_data.package_data["multifilepkg-stubs"]) == 7
 
 
 def test_uploaded_packages() -> None: