xref: /aosp_15_r20/external/pigweed/pw_build/py/pw_build/mirror_tree.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"""Mirrors a directory tree to another directory using hard links."""
15
16import argparse
17import os
18from pathlib import Path
19import shutil
20from typing import Iterable, Iterator
21
22
23def _parse_args() -> argparse.Namespace:
24    """Registers the script's arguments on an argument parser."""
25
26    parser = argparse.ArgumentParser(description=__doc__)
27
28    parser.add_argument(
29        '--source-root',
30        type=Path,
31        required=True,
32        help='Prefix to strip from the source files',
33    )
34    parser.add_argument(
35        'sources', type=Path, nargs='*', help='Files to mirror to the directory'
36    )
37    parser.add_argument(
38        '--directory',
39        type=Path,
40        required=True,
41        help='Directory to which to mirror the sources',
42    )
43    parser.add_argument(
44        '--path-file', type=Path, help='File with paths to files to mirror'
45    )
46
47    return parser.parse_args()
48
49
50def _link_files(
51    source_root: Path, sources: Iterable[Path], directory: Path
52) -> Iterator[Path]:
53    for source in sources:
54        dest = directory / source.relative_to(source_root)
55        dest.parent.mkdir(parents=True, exist_ok=True)
56
57        if dest.exists():
58            dest.unlink()
59
60        # Use a hard link to avoid unnecessary copies. Resolve the source before
61        # linking in case it is a symlink.
62        source = source.resolve()
63        try:
64            os.link(source, dest)
65            yield dest
66
67        # If the link failed try copying. If copying fails re-raise the
68        # original exception.
69        except OSError:
70            shutil.copy(source, dest)
71            yield dest
72
73
74def _link_files_or_dirs(
75    paths: Iterable[Path], directory: Path
76) -> Iterator[Path]:
77    """Links files or directories into the output directory.
78
79    Files are linked directly; files in directories are linked as relative paths
80    from the directory.
81    """
82
83    for path in paths:
84        if path.is_dir():
85            files = (p for p in path.glob('**/*') if p.is_file())
86            yield from _link_files(path, files, directory)
87        elif path.is_file():
88            yield from _link_files(path.parent, [path], directory)
89        else:
90            raise FileNotFoundError(f'{path} does not exist!')
91
92
93def mirror_paths(
94    source_root: Path,
95    sources: Iterable[Path],
96    directory: Path,
97    path_file: Path | None = None,
98) -> list[Path]:
99    """Creates hard links in the provided directory for the provided sources.
100
101    Args:
102      source_root: Base path for files in sources.
103      sources: Files to link to from the directory.
104      directory: The output directory.
105      path_file: A file with file or directory paths to link to.
106    """
107    directory.mkdir(parents=True, exist_ok=True)
108
109    outputs = list(_link_files(source_root, sources, directory))
110
111    if path_file:
112        paths = (Path(p).resolve() for p in path_file.read_text().splitlines())
113        outputs.extend(_link_files_or_dirs(paths, directory))
114
115    return outputs
116
117
118if __name__ == '__main__':
119    mirror_paths(**vars(_parse_args()))
120