xref: /aosp_15_r20/external/pigweed/pw_build/py/pw_build/generate_python_package.py (revision 61c4878ac05f98d0ceed94b57d316916de578985)
1# Copyright 2021 The Pigweed Authors
2#
3# Licensed under the Apache License, Version 2.0 (the "License"); you may not
4# use this file except in compliance with the License. You may obtain a copy of
5# the License at
6#
7#     https://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
11# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
12# License for the specific language governing permissions and limitations under
13# the License.
14"""Generates a setup.cfg file for a pw_python_package target."""
15
16import argparse
17from collections import defaultdict
18import configparser
19from dataclasses import dataclass
20import json
21from pathlib import Path
22import sys
23import textwrap
24from typing import (
25    Any,
26    Iterable,
27    Iterator,
28    Sequence,
29    Set,
30    TextIO,
31)
32
33try:
34    from pw_build.mirror_tree import mirror_paths
35except ImportError:
36    # Append this path to the module search path to allow running this module
37    # before the pw_build package is installed.
38    sys.path.append(str(Path(__file__).resolve().parent.parent))
39    from pw_build.mirror_tree import mirror_paths
40
41
42def _parse_args() -> argparse.Namespace:
43    parser = argparse.ArgumentParser(description=__doc__)
44
45    parser.add_argument('--label', help='Label for this Python package')
46    parser.add_argument(
47        '--proto-library',
48        dest='proto_libraries',
49        type=argparse.FileType('r'),
50        default=[],
51        action='append',
52        help='Paths',
53    )
54    parser.add_argument(
55        '--generated-root',
56        required=True,
57        type=Path,
58        help='The base directory for the Python package',
59    )
60    parser.add_argument(
61        '--setup-json',
62        required=True,
63        type=argparse.FileType('r'),
64        help='setup.py keywords as JSON',
65    )
66    parser.add_argument(
67        '--module-as-package',
68        action='store_true',
69        help='Generate an __init__.py that imports everything',
70    )
71    parser.add_argument(
72        'files',
73        type=Path,
74        nargs='+',
75        help='Relative paths to the files in the package',
76    )
77    return parser.parse_args()
78
79
80def _check_nested_protos(label: str, proto_info: dict[str, Any]) -> None:
81    """Checks that the proto library refers to this package."""
82    python_package = proto_info['nested_in_python_package']
83
84    if python_package != label:
85        raise ValueError(
86            f"{label}'s 'proto_library' is set to {proto_info['label']}, but "
87            f"that target's 'python_package' is {python_package or 'not set'}. "
88            f"Set {proto_info['label']}'s 'python_package' to {label}."
89        )
90
91
92@dataclass(frozen=True)
93class _ProtoInfo:
94    root: Path
95    sources: Sequence[Path]
96    deps: Sequence[str]
97
98
99def _collect_all_files(
100    root: Path, files: list[Path], paths_to_collect: Iterable[_ProtoInfo]
101) -> dict[str, Set[str]]:
102    """Collects files in output dir, adds to files; returns package_data."""
103    root.mkdir(exist_ok=True)
104
105    for proto_info in paths_to_collect:
106        # Mirror the proto files to this package.
107        files += mirror_paths(proto_info.root, proto_info.sources, root)
108
109    # Find all subpackages, including empty ones.
110    subpackages: Set[Path] = set()
111    for file in (f.relative_to(root) for f in files):
112        subpackages.update(root / path for path in file.parents)
113    subpackages.remove(root)
114
115    # Make sure there are __init__.py and py.typed files for each subpackage.
116    for pkg in subpackages:
117        pytyped = pkg / 'py.typed'
118        if not pytyped.exists():
119            pytyped.touch()
120        files.append(pytyped)
121
122        # Create an __init__.py file if it doesn't already exist.
123        initpy = pkg / '__init__.py'
124        if not initpy.exists():
125            # Use pkgutil.extend_path to treat this as a namespaced package.
126            # This allows imports with the same name to live in multiple
127            # separate PYTHONPATH locations.
128            initpy.write_text(
129                'from pkgutil import extend_path  # type: ignore\n'
130                '__path__ = extend_path(__path__, __name__)  # type: ignore\n'
131            )
132        files.append(initpy)
133
134    pkg_data: dict[str, Set[str]] = defaultdict(set)
135
136    # Add all non-source files to package data.
137    for file in (f for f in files if f.suffix != '.py'):
138        pkg = file.parent
139
140        package_name = pkg.relative_to(root).as_posix().replace('/', '.')
141        pkg_data[package_name].add(file.name)
142
143    return pkg_data
144
145
146PYPROJECT_FILE = '''\
147# Generated file. Do not modify.
148[build-system]
149requires = ['setuptools', 'wheel']
150build-backend = 'setuptools.build_meta'
151'''
152
153
154def _get_setup_keywords(pkg_data: dict, keywords: dict) -> dict:
155    """Gather all setuptools.setup() keyword args."""
156    options_keywords = dict(
157        packages=list(pkg_data),
158        package_data={pkg: list(files) for pkg, files in pkg_data.items()},
159    )
160
161    keywords['options'].update(options_keywords)
162    return keywords
163
164
165def _write_to_config(
166    config: configparser.ConfigParser, data: dict, section: str | None = None
167):
168    """Populate a ConfigParser instance with the contents of a dict."""
169    # Add a specified section if missing.
170    if section is not None and not config.has_section(section):
171        config.add_section(section)
172
173    for key, value in data.items():
174        # Value is a dict so create a new subsection
175        if isinstance(value, dict):
176            _write_to_config(
177                config,
178                value,
179                f'{section}.{key}' if section else key,
180            )
181        elif isinstance(value, list):
182            if value:
183                assert section is not None
184                # Convert the list to an allowed str format.
185                config[section][key] = '\n' + '\n'.join(value)
186        else:
187            assert section is not None
188            # Add the value as a string. See expected types here:
189            # https://setuptools.readthedocs.io/en/latest/userguide/declarative_config.html#specifying-values
190            config[section][key] = str(value)
191
192
193def _generate_setup_cfg(
194    pkg_data: dict,
195    keywords: dict,
196    config_file_path: Path,
197) -> None:
198    """Creates a setup.cfg file based on setuptools keywords."""
199    setup_keywords = _get_setup_keywords(pkg_data, keywords)
200
201    config = configparser.ConfigParser()
202
203    _write_to_config(config, setup_keywords)
204
205    # Write the config to a file.
206    with config_file_path.open('w') as config_file:
207        config.write(config_file)
208
209
210def _import_module_in_package_init(all_files: list[Path]) -> None:
211    """Generates an __init__.py that imports the module.
212
213    This makes an individual module usable as a package. This is used for proto
214    modules.
215    """
216    sources = [
217        f for f in all_files if f.suffix == '.py' and f.name != '__init__.py'
218    ]
219    assert (
220        len(sources) == 1
221    ), 'Module as package expects a single .py source file'
222
223    (source,) = sources
224    source.parent.joinpath('__init__.py').write_text(
225        f'from {source.stem}.{source.stem} import *\n'
226    )
227
228
229def _load_metadata(
230    label: str, proto_libraries: Iterable[TextIO]
231) -> Iterator[_ProtoInfo]:
232    for proto_library_file in proto_libraries:
233        info = json.load(proto_library_file)
234        _check_nested_protos(label, info)
235
236        deps = []
237        for dep in info['dependencies']:
238            with open(dep) as file:
239                deps.append(json.load(file)['package'])
240
241        yield _ProtoInfo(
242            Path(info['root']),
243            tuple(Path(p) for p in info['protoc_outputs']),
244            deps,
245        )
246
247
248def main(
249    generated_root: Path,
250    files: list[Path],
251    module_as_package: bool,
252    setup_json: TextIO,
253    label: str,
254    proto_libraries: Iterable[TextIO],
255) -> int:
256    """Generates a setup.py and other files for a Python package."""
257    proto_infos = list(_load_metadata(label, proto_libraries))
258    try:
259        pkg_data = _collect_all_files(generated_root, files, proto_infos)
260    except ValueError as error:
261        msg = '\n'.join(textwrap.wrap(str(error), 78))
262        print(
263            f'ERROR: Failed to generate Python package {label}:\n\n'
264            f'{textwrap.indent(msg, "  ")}\n',
265            file=sys.stderr,
266        )
267        return 1
268
269    with setup_json:
270        setup_keywords = json.load(setup_json)
271        setup_keywords.setdefault('options', {})
272
273    setup_keywords['options'].setdefault('install_requires', [])
274
275    if module_as_package:
276        _import_module_in_package_init(files)
277
278    # Create the pyproject.toml and setup.cfg files for this package.
279    generated_root.joinpath('pyproject.toml').write_text(PYPROJECT_FILE)
280    _generate_setup_cfg(
281        pkg_data,
282        setup_keywords,
283        config_file_path=generated_root.joinpath('setup.cfg'),
284    )
285
286    return 0
287
288
289if __name__ == '__main__':
290    sys.exit(main(**vars(_parse_args())))
291