1# Copyright (c) 2024 Chen Xingyu <hi@xingrz.me>
2# SPDX-License-Identifier: Apache-2.0
3
4'''Runner for probe-rs.'''
5
6from runners.core import RunnerCaps, ZephyrBinaryRunner
7
8
9class ProbeRsBinaryRunner(ZephyrBinaryRunner):
10    '''Runner front-end for probe-rs.'''
11
12    def __init__(self, cfg, chip,
13                 probe_rs='probe-rs',
14                 dev_id=None,
15                 erase=False,
16                 tool_opt=None):
17        super().__init__(cfg)
18
19        self.probe_rs = probe_rs
20        self.erase = erase
21
22        self.args = ['--chip', chip]
23
24        if dev_id is not None:
25            self.args += ['--probe', dev_id]
26
27        if tool_opt is not None:
28            self.args += tool_opt
29
30        self.elf_name = cfg.elf_file
31
32    @classmethod
33    def name(cls):
34        return 'probe-rs'
35
36    @classmethod
37    def capabilities(cls):
38        return RunnerCaps(commands={'flash'},
39                          dev_id=True,
40                          erase=True,
41                          tool_opt=True)
42
43    @classmethod
44    def do_add_parser(cls, parser):
45        parser.add_argument('--chip', required=True,
46                            help='chip name')
47        parser.add_argument('--probe-rs', default='probe-rs',
48                            help='path to probe-rs tool, default is probe-rs')
49
50    @classmethod
51    def dev_id_help(cls) -> str:
52        return '''select a specific probe, in the form `VID:PID:<Serial>`'''
53
54    @classmethod
55    def tool_opt_help(cls) -> str:
56        return '''additional options for probe-rs,
57                  e.g. --chip-description-path=/path/to/chip.yml'''
58
59    @classmethod
60    def do_create(cls, cfg, args):
61        return ProbeRsBinaryRunner(cfg, args.chip,
62                                   probe_rs=args.probe_rs,
63                                   dev_id=args.dev_id,
64                                   erase=args.erase,
65                                   tool_opt=args.tool_opt)
66
67    def do_run(self, command, **kwargs):
68        self.require(self.probe_rs)
69        if command == 'flash':
70            self.do_flash(**kwargs)
71
72    def do_flash(self, **kwargs):
73        download_args = []
74        if self.erase:
75            download_args += ['--chip-erase']
76        download_args += [self.elf_name]
77
78        self.check_call([self.probe_rs, 'download']
79                        + self.args + download_args)
80
81        self.check_call([self.probe_rs, 'reset']
82                        + self.args)
83