xref: /aosp_15_r20/external/pigweed/pw_module/py/pw_module/seed.py (revision 61c4878ac05f98d0ceed94b57d316916de578985)
1# Copyright 2024 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"""Manages SEED documents in Pigweed."""
15
16import argparse
17import datetime
18import enum
19from dataclasses import dataclass
20from pathlib import Path
21import random
22import re
23import subprocess
24import sys
25import urllib.request
26from typing import Dict, Iterable, List, Optional, Tuple, Union
27
28import pw_cli.color
29import pw_cli.env
30from pw_cli.git_repo import GitRepo
31from pw_cli.tool_runner import BasicSubprocessRunner
32
33_NEW_SEED_TEMPLATE = '''.. _seed-{num:04d}:
34
35{title_underline}
36{formatted_title}
37{title_underline}
38.. seed::
39   :number: {num}
40   :name: {title}
41   :status: {status}
42   :proposal_date: {date}
43   :cl: {changelist}
44   :authors: {authors}
45   :facilitator: Unassigned
46
47-------
48Summary
49-------
50Write up your proposal here.
51'''
52
53
54class SeedStatus(enum.Enum):
55    """Possible states of a SEED proposal."""
56
57    DRAFT = 0
58    OPEN_FOR_COMMENTS = 1
59    LAST_CALL = 2
60    ACCEPTED = 3
61    REJECTED = 4
62    DEPRECATED = 5
63    SUPERSEDED = 6
64    ON_HOLD = 7
65    META = 8
66
67    def __str__(self) -> str:
68        if self is SeedStatus.DRAFT:
69            return 'Draft'
70        if self is SeedStatus.OPEN_FOR_COMMENTS:
71            return 'Open for Comments'
72        if self is SeedStatus.LAST_CALL:
73            return 'Last Call'
74        if self is SeedStatus.ACCEPTED:
75            return 'Accepted'
76        if self is SeedStatus.REJECTED:
77            return 'Rejected'
78        if self is SeedStatus.DEPRECATED:
79            return 'Deprecated'
80        if self is SeedStatus.SUPERSEDED:
81            return 'Superseded'
82        if self is SeedStatus.ON_HOLD:
83            return 'On Hold'
84        if self is SeedStatus.META:
85            return 'Meta'
86
87        return ''
88
89
90@dataclass
91class SeedMetadata:
92    number: int
93    title: str
94    authors: str
95    status: SeedStatus
96    changelist: Optional[int] = None
97    sources: Optional[List[str]] = None
98
99    def default_filename(self) -> str:
100        return f'{self.number:04d}.rst'
101
102
103class SeedRegistry:
104    """
105    Represents a registry of SEEDs located somewhere in pigweed.git, which can
106    be read, modified, and written.
107
108    Currently, this is implemented as a basic text parser for the BUILD.gn file
109    in the seed/ directory; however, in the future it may be rewritten to use a
110    different backing data source.
111    """
112
113    class _State(enum.Enum):
114        OUTER = 0
115        SEED = 1
116        INDEX = 2
117        INDEX_SEEDS_LIST = 3
118
119    @classmethod
120    def parse(cls, registry_file: Path) -> 'SeedRegistry':
121        return cls(registry_file)
122
123    def __init__(self, seed_build_file: Path):
124        self._file = seed_build_file
125        self._lines = seed_build_file.read_text().split('\n')
126        self._seeds: Dict[int, Tuple[int, int]] = {}
127        self._next_seed_number = 101
128
129        seed_regex = re.compile(r'pw_seed\("(\d+)"\)')
130
131        state = SeedRegistry._State.OUTER
132
133        section_start_index = 0
134        current_seed_number = 0
135
136        # Run through the GN file, doing some basic parsing of its targets.
137        for i, line in enumerate(self._lines):
138            if state is SeedRegistry._State.OUTER:
139                seed_match = seed_regex.search(line)
140                if seed_match:
141                    # SEED definition target: extract the number.
142                    state = SeedRegistry._State.SEED
143
144                    section_start_index = i
145                    current_seed_number = int(seed_match.group(1))
146
147                    if current_seed_number >= self._next_seed_number:
148                        self._next_seed_number = current_seed_number + 1
149
150                if line == 'pw_seed_index("seeds") {':
151                    state = SeedRegistry._State.INDEX
152                    # Skip back past the comments preceding the SEED index
153                    # target. New SEEDs will be inserted here.
154                    insertion_index = i
155                    while insertion_index > 0 and self._lines[
156                        insertion_index - 1
157                    ].startswith('#'):
158                        insertion_index -= 1
159
160                    self._seed_insertion_index = insertion_index
161
162            if state is SeedRegistry._State.SEED:
163                if line == '}':
164                    self._seeds[current_seed_number] = (section_start_index, i)
165                    state = SeedRegistry._State.OUTER
166
167            if state is SeedRegistry._State.INDEX:
168                if line == '}':
169                    state = SeedRegistry._State.OUTER
170                if line == '  seeds = [':
171                    state = SeedRegistry._State.INDEX_SEEDS_LIST
172
173            if state is SeedRegistry._State.INDEX_SEEDS_LIST:
174                if line == '  ]':
175                    self._index_seeds_end = i
176                    state = SeedRegistry._State.INDEX
177
178    def file(self) -> Path:
179        """Returns the file which backs this registry."""
180        return self._file
181
182    def seed_count(self) -> int:
183        """Returns the number of SEEDs registered."""
184        return len(self._seeds)
185
186    def insert(self, seed: SeedMetadata) -> None:
187        """Adds a new seed to the registry."""
188
189        new_seed = [
190            f'pw_seed("{seed.number:04d}") {{',
191            f'  title = "{seed.title}"',
192            f'  author = "{seed.authors}"',
193            f'  status = "{seed.status}"',
194        ]
195
196        if seed.changelist is not None:
197            new_seed.append(f'  changelist = {seed.changelist}')
198
199        if seed.sources is not None:
200            if len(seed.sources) == 0:
201                new_seed.append('  sources = []')
202            elif len(seed.sources) == 1:
203                new_seed.append(f'  sources = [ "{seed.sources[0]}" ]')
204            else:
205                new_seed.append('  sources = [')
206                new_seed.extend(f'    "{source}",' for source in seed.sources)
207                new_seed.append('  ]')
208
209        new_seed += [
210            '}',
211            '',
212        ]
213        self._lines = (
214            self._lines[: self._seed_insertion_index]
215            + new_seed
216            + self._lines[self._seed_insertion_index : self._index_seeds_end]
217            + [f'    ":{seed.number:04d}",']
218            + self._lines[self._index_seeds_end :]
219        )
220
221        self._seed_insertion_index += len(new_seed)
222        self._index_seeds_end += len(new_seed)
223
224        if seed.number == self._next_seed_number:
225            self._next_seed_number += 1
226
227    def next_seed_number(self) -> int:
228        return self._next_seed_number
229
230    def write(self) -> None:
231        self._file.write_text('\n'.join(self._lines))
232
233
234_GERRIT_HOOK_URL = (
235    'https://gerrit-review.googlesource.com/tools/hooks/commit-msg'
236)
237
238
239# TODO: pwbug.dev/318746837 - Extract this to somewhere more general.
240def _install_gerrit_hook(git_root: Path) -> None:
241    hook_file = git_root / '.git' / 'hooks' / 'commit-msg'
242    urllib.request.urlretrieve(_GERRIT_HOOK_URL, hook_file)
243    hook_file.chmod(0o755)
244
245
246def _request_new_seed_metadata(
247    repo: GitRepo,
248    registry: SeedRegistry,
249    colors,
250) -> SeedMetadata:
251    if repo.has_uncommitted_changes():
252        print(
253            colors.red(
254                'You have uncommitted Git changes. '
255                'Please commit or stash before creating a SEED.'
256            )
257        )
258        sys.exit(1)
259
260    print(
261        colors.yellow(
262            'This command will create Git commits. '
263            'Make sure you are on a suitable branch.'
264        )
265    )
266    print(f'Current branch: {colors.cyan(repo.current_branch())}')
267    print('')
268
269    number = registry.next_seed_number()
270
271    try:
272        num = input(
273            f'SEED number (default={colors.bold_white(number)}): '
274        ).strip()
275
276        while True:
277            try:
278                if num:
279                    number = int(num)
280                break
281            except ValueError:
282                num = input('Invalid number entered. Try again: ').strip()
283
284        title = input('SEED title: ').strip()
285        while not title:
286            title = input(
287                'Title cannot be empty. Re-enter SEED title: '
288            ).strip()
289
290        authors = input('SEED authors: ').strip()
291        while not authors:
292            authors = input(
293                'Authors list cannot be empty. Re-enter SEED authors: '
294            ).strip()
295
296        print('The following SEED will be created.')
297        print('')
298        print(f'  Number: {colors.green(number)}')
299        print(f'  Title: {colors.green(title)}')
300        print(f'  Authors: {colors.green(authors)}')
301        print('')
302        print(
303            'This will create two commits on branch '
304            + colors.cyan(repo.current_branch())
305            + ' and push them to Gerrit.'
306        )
307
308        create = True
309        confirm = input(f'Proceed? [{colors.bold_white("Y")}/n] ').strip()
310        if confirm:
311            create = confirm == 'Y'
312
313    except KeyboardInterrupt:
314        print('\nReceived CTRL-C, exiting...')
315        sys.exit(0)
316
317    if not create:
318        sys.exit(0)
319
320    return SeedMetadata(
321        number=number,
322        title=title,
323        authors=authors,
324        status=SeedStatus.DRAFT,
325    )
326
327
328@dataclass
329class GerritChange:
330    id: str
331    number: int
332    title: str
333
334    def url(self) -> str:
335        return (
336            'https://pigweed-review.googlesource.com'
337            f'/c/pigweed/pigweed/+/{self.number}'
338        )
339
340
341def commit_and_push(
342    repo: GitRepo,
343    files: Iterable[Union[Path, str]],
344    commit_message: str,
345    change_id: Optional[str] = None,
346) -> GerritChange:
347    """Creates a commit with the given files and message and pushes to Gerrit.
348
349    Args:
350        change_id: Optional Gerrit change ID to use. If not specified, generates
351            a new one.
352    """
353    if change_id is not None:
354        commit_message = f'{commit_message}\n\nChange-Id: {change_id}'
355
356    subprocess.run(
357        ['git', 'add'] + list(files), capture_output=True, check=True
358    )
359    subprocess.run(
360        ['git', 'commit', '-m', commit_message], capture_output=True, check=True
361    )
362
363    if change_id is None:
364        # Parse the generated change ID from the commit if it wasn't
365        # explicitly set.
366        change_id = repo.commit_change_id()
367
368        if change_id is None:
369            # If the commit doesn't have a Change-Id, the Gerrit hook is not
370            # installed. Install it and try modifying the commit.
371            _install_gerrit_hook(repo.root())
372            subprocess.run(
373                ['git', 'commit', '--amend', '--no-edit'],
374                capture_output=True,
375                check=True,
376            )
377            change_id = repo.commit_change_id()
378            assert change_id is not None
379
380    process = subprocess.run(
381        [
382            'git',
383            'push',
384            'origin',
385            '+HEAD:refs/for/main',
386            '--no-verify',
387        ],
388        capture_output=True,
389        text=True,
390        check=True,
391    )
392
393    output = process.stderr
394
395    regex = re.compile(
396        '^\\s*remote:\\s*'
397        'https://pigweed-review.(?:git.corp.google|googlesource).com/'
398        'c/pigweed/pigweed/\\+/(?P<num>\\d+)\\s+',
399        re.MULTILINE,
400    )
401    match = regex.search(output)
402    if not match:
403        raise ValueError(f"invalid output from 'git push': {output}")
404    change_num = int(match.group('num'))
405
406    return GerritChange(change_id, change_num, commit_message.split('\n')[0])
407
408
409def _create_wip_seed_doc_change(
410    repo: GitRepo,
411    new_seed: SeedMetadata,
412) -> GerritChange:
413    """Commits and pushes a boilerplate CL doc to Gerrit.
414
415    Returns information about the CL.
416    """
417    env = pw_cli.env.pigweed_environment()
418    seed_rst_file = env.PW_ROOT / 'seed' / new_seed.default_filename()
419
420    formatted_title = f'{new_seed.number:04d}: {new_seed.title}'
421    title_underline = '=' * len(formatted_title)
422
423    seed_rst_file.write_text(
424        _NEW_SEED_TEMPLATE.format(
425            formatted_title=formatted_title,
426            title_underline=title_underline,
427            num=new_seed.number,
428            title=new_seed.title,
429            authors=new_seed.authors,
430            status=new_seed.status,
431            date=datetime.date.today().strftime('%Y-%m-%d'),
432            changelist=0,
433        )
434    )
435
436    temp_branch = f'wip-seed-{new_seed.number}-{random.randrange(0, 2**16):x}'
437    subprocess.run(
438        ['git', 'checkout', '-b', temp_branch], capture_output=True, check=True
439    )
440
441    commit_message = f'SEED-{new_seed.number:04d}: {new_seed.title}'
442
443    try:
444        cl = commit_and_push(repo, [seed_rst_file], commit_message)
445    except subprocess.CalledProcessError as err:
446        print(f'Command {err.cmd} failed; stderr:')
447        print(err.stderr)
448        sys.exit(1)
449    finally:
450        subprocess.run(
451            ['git', 'checkout', '-'], capture_output=True, check=True
452        )
453        subprocess.run(
454            ['git', 'branch', '-D', temp_branch],
455            capture_output=True,
456            check=True,
457        )
458
459    return cl
460
461
462def create_seed_number_claim_change(
463    repo: GitRepo,
464    new_seed: SeedMetadata,
465    registry: SeedRegistry,
466) -> GerritChange:
467    commit_message = f'SEED-{new_seed.number:04d}: Claim SEED number'
468    registry.insert(new_seed)
469    registry.write()
470    return commit_and_push(repo, [registry.file()], commit_message)
471
472
473def _create_seed_doc_change(
474    repo: GitRepo,
475    new_seed: SeedMetadata,
476    registry: SeedRegistry,
477    wip_change: GerritChange,
478) -> None:
479    env = pw_cli.env.pigweed_environment()
480    seed_rst_file = env.PW_ROOT / 'seed' / new_seed.default_filename()
481
482    formatted_title = f'{new_seed.number:04d}: {new_seed.title}'
483    title_underline = '=' * len(formatted_title)
484
485    seed_rst_file.write_text(
486        _NEW_SEED_TEMPLATE.format(
487            formatted_title=formatted_title,
488            title_underline=title_underline,
489            num=new_seed.number,
490            title=new_seed.title,
491            authors=new_seed.authors,
492            status=new_seed.status,
493            date=datetime.date.today().strftime('%Y-%m-%d'),
494            changelist=new_seed.changelist,
495        )
496    )
497
498    new_seed.sources = [seed_rst_file.relative_to(registry.file().parent)]
499    new_seed.changelist = None
500    registry.insert(new_seed)
501    registry.write()
502
503    commit_message = f'SEED-{new_seed.number:04d}: {new_seed.title}'
504    commit_and_push(
505        repo,
506        [registry.file(), seed_rst_file],
507        commit_message,
508        change_id=wip_change.id,
509    )
510
511
512def create_seed() -> int:
513    colors = pw_cli.color.colors()
514    env = pw_cli.env.pigweed_environment()
515
516    repo = GitRepo(env.PW_ROOT, BasicSubprocessRunner())
517
518    registry_path = env.PW_ROOT / 'seed' / 'BUILD.gn'
519
520    wip_registry = SeedRegistry.parse(registry_path)
521    registry = SeedRegistry.parse(registry_path)
522
523    seed = _request_new_seed_metadata(repo, wip_registry, colors)
524    seed_cl = _create_wip_seed_doc_change(repo, seed)
525
526    seed.changelist = seed_cl.number
527
528    number_cl = create_seed_number_claim_change(repo, seed, wip_registry)
529    _create_seed_doc_change(repo, seed, registry, seed_cl)
530
531    print()
532    print(f'Created two CLs for SEED-{seed.number:04d}:')
533    print()
534    print(f'-  {number_cl.title}')
535    print(f'   <{number_cl.url()}>')
536    print()
537    print(f'-  {seed_cl.title}')
538    print(f'   <{seed_cl.url()}>')
539
540    return 0
541
542
543def parse_args() -> argparse.Namespace:
544    parser = argparse.ArgumentParser(description=__doc__)
545    parser.set_defaults(func=lambda **_kwargs: parser.print_help())
546
547    subparsers = parser.add_subparsers(title='subcommands')
548
549    # No args for now, as this initially only runs in interactive mode.
550    create_parser = subparsers.add_parser('create', help='Creates a new SEED')
551    create_parser.set_defaults(func=create_seed)
552
553    return parser.parse_args()
554
555
556def main() -> int:
557    args = {**vars(parse_args())}
558    func = args['func']
559    del args['func']
560
561    exit_code = func(**args)
562    return 0 if exit_code is None else exit_code
563
564
565if __name__ == '__main__':
566    sys.exit(main())
567