xref: /aosp_15_r20/external/pigweed/pw_build/py/pw_build/collect_wheels.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"""Collect Python wheels from a build into a central directory."""
15
16import argparse
17import logging
18from pathlib import Path
19import shutil
20import sys
21
22_LOG = logging.getLogger(__name__)
23
24
25def _parse_args():
26    parser = argparse.ArgumentParser(description=__doc__)
27    parser.add_argument(
28        '--prefix',
29        type=Path,
30        help='Root search path to use in conjunction with --wheels_file',
31    )
32    parser.add_argument(
33        '--suffix-file',
34        type=Path,
35        help=(
36            'File that lists subdirs relative to --prefix, one per line,'
37            'to search for .whl files to copy into --out_dir'
38        ),
39    )
40    parser.add_argument(
41        '--out-dir',
42        type=Path,
43        help='Path where all the built and collected .whl files should be put',
44    )
45
46    return parser.parse_args()
47
48
49def copy_wheels(prefix: Path, suffix_file: Path, out_dir: Path) -> None:
50    """Copy Python wheels or source archives to the out_dir.
51
52    Raises:
53      FileExistsError: If any separate wheel files are copied to the same
54          destination file path.
55    """
56    # Delete existing wheels from the out dir, there may be stale versions.
57    shutil.rmtree(out_dir, ignore_errors=True)
58    out_dir.mkdir(parents=True, exist_ok=True)
59
60    copied_files: dict[str, Path] = dict()
61    requirements_content: str = ''
62
63    for suffix in suffix_file.read_text().splitlines():
64        path = prefix / suffix.strip()
65        _LOG.debug('Searching for wheels in %s', path)
66        if path == out_dir:
67            continue
68        for wheel in path.iterdir():
69            if wheel.suffix not in ('.gz', '.whl'):
70                continue
71
72            if wheel.name in copied_files:
73                _LOG.error(
74                    'Attempting to override %s with %s',
75                    copied_files[wheel.name],
76                    wheel,
77                )
78                raise FileExistsError(
79                    f'{wheel.name} conflict: '
80                    f'{copied_files[wheel.name]} and {wheel}'
81                )
82            copied_files[wheel.name] = wheel
83            _LOG.debug('Copying %s to %s', wheel, out_dir)
84            shutil.copy(wheel, out_dir)
85            requirements_file = wheel.parent / 'requirements.txt'
86
87            if requirements_file.is_file():
88                requirements_content += '\n'
89                requirements_content += requirements_file.read_text()
90
91    if requirements_content:
92        (out_dir / 'requirements.txt').write_text(
93            '# Auto-generated requirements.txt\n' + requirements_content,
94            encoding='utf-8',
95        )
96
97
98def main() -> None:
99    copy_wheels(**vars(_parse_args()))
100
101
102if __name__ == '__main__':
103    logging.basicConfig()
104    main()
105    sys.exit(0)
106