1#!/usr/bin/env python3
2
3# Copyright (c) 2024 Vestas Wind Systems A/S
4# Copyright (c) 2020 Nordic Semiconductor ASA
5# SPDX-License-Identifier: Apache-2.0
6
7import argparse
8import json
9import sys
10from dataclasses import dataclass
11from pathlib import Path
12
13import pykwalify.core
14import yaml
15
16try:
17    from yaml import CSafeLoader as SafeLoader
18except ImportError:
19    from yaml import SafeLoader
20
21SHIELD_SCHEMA_PATH = str(Path(__file__).parent / 'schemas' / 'shield-schema.yml')
22with open(SHIELD_SCHEMA_PATH) as f:
23    shield_schema = yaml.load(f.read(), Loader=SafeLoader)
24
25SHIELD_YML = 'shield.yml'
26
27#
28# This is shared code between the build system's 'shields' target
29# and the 'west shields' extension command. If you change it, make
30# sure to test both ways it can be used.
31#
32# (It's done this way to keep west optional, making it possible to run
33# 'ninja shields' in a build directory without west installed.)
34#
35
36@dataclass(frozen=True)
37class Shield:
38    name: str
39    dir: Path
40    full_name: str | None = None
41    vendor: str | None = None
42    supported_features: list[str] | None = None
43
44def shield_key(shield):
45    return shield.name
46
47def process_shield_data(shield_data, shield_dir):
48    # Create shield from yaml data
49    return Shield(
50        name=shield_data['name'],
51        dir=shield_dir,
52        full_name=shield_data.get('full_name'),
53        vendor=shield_data.get('vendor'),
54        supported_features=shield_data.get('supported_features', []),
55    )
56
57def find_shields(args):
58    ret = []
59
60    for root in args.board_roots:
61        for shields in find_shields_in(root):
62            ret.append(shields)
63
64    return sorted(ret, key=shield_key)
65
66def find_shields_in(root):
67    shields = root / 'boards' / 'shields'
68    ret = []
69
70    if not shields.exists():
71        return ret
72
73    for maybe_shield in (shields).iterdir():
74        if not maybe_shield.is_dir():
75            continue
76
77        # Check for shield.yml first
78        shield_yml = maybe_shield / SHIELD_YML
79        if shield_yml.is_file():
80            with shield_yml.open('r', encoding='utf-8') as f:
81                shield_data = yaml.load(f.read(), Loader=SafeLoader)
82
83            try:
84                pykwalify.core.Core(source_data=shield_data, schema_data=shield_schema).validate()
85            except pykwalify.errors.SchemaError as e:
86                sys.exit(f'ERROR: Malformed shield.yml in file: {shield_yml.as_posix()}\n{e}')
87
88            if 'shields' in shield_data:
89                # Multiple shields format
90                for shield_info in shield_data['shields']:
91                    ret.append(process_shield_data(shield_info, maybe_shield))
92            elif 'shield' in shield_data:
93                # Single shield format
94                ret.append(process_shield_data(shield_data['shield'], maybe_shield))
95            continue
96
97        # Fallback to legacy method if no shield.yml
98        for maybe_kconfig in maybe_shield.iterdir():
99            if maybe_kconfig.name == 'Kconfig.shield':
100                for maybe_overlay in maybe_shield.iterdir():
101                    file_name = maybe_overlay.name
102                    if file_name.endswith('.overlay'):
103                        shield_name = file_name[:-len('.overlay')]
104                        ret.append(Shield(shield_name, maybe_shield))
105
106    return sorted(ret, key=shield_key)
107
108def parse_args():
109    parser = argparse.ArgumentParser(allow_abbrev=False)
110    add_args(parser)
111    add_args_formatting(parser)
112    return parser.parse_args()
113
114def add_args(parser):
115    # Remember to update west-completion.bash if you add or remove
116    # flags
117    parser.add_argument("--board-root", dest='board_roots', default=[],
118                        type=Path, action='append',
119                        help='add a board root, may be given more than once')
120
121def add_args_formatting(parser):
122    parser.add_argument("--json", action='store_true',
123                        help='''output list of shields in JSON format''')
124
125def dump_shields(shields):
126    if args.json:
127        print(
128            json.dumps([{'dir': shield.dir.as_posix(), 'name': shield.name} for shield in shields])
129        )
130    else:
131        for shield in shields:
132            print(f'  {shield.name}')
133
134if __name__ == '__main__':
135    args = parse_args()
136    dump_shields(find_shields(args))
137