1- # Copyright 2024 ARM Limited
1+ # Copyright 2024-2025 ARM Limited
22#
33# Licensed under the Apache License, Version 2.0 (the "License");
44# you may not use this file except in compliance with the License.
2020import logging
2121import os
2222import time
23+
2324from platform import machine
25+ from typing import Optional , cast , Protocol , TYPE_CHECKING , TypedDict , Union
26+ from typing_extensions import NotRequired , LiteralString
27+ if TYPE_CHECKING :
28+ from _typeshed import StrPath , BytesPath
29+ from devlib .platform import Platform
30+ else :
31+ StrPath = str
32+ BytesPath = bytes
2433
2534from devlib .exception import (TargetStableError , HostError )
26- from devlib .target import LinuxTarget
35+ from devlib .target import LinuxTarget , Target
2736from devlib .utils .misc import get_subprocess , which
2837from devlib .utils .ssh import SshConnection
38+ from devlib .utils .annotation_helpers import SubprocessCommand , SshUserConnectionSettings
2939
3040
3141class TargetRunner :
@@ -40,15 +50,22 @@ class TargetRunner:
4050 """
4151
4252 def __init__ (self ,
43- target ) :
53+ target : Target ) -> None :
4454 self .target = target
45-
4655 self .logger = logging .getLogger (self .__class__ .__name__ )
4756
4857 def __enter__ (self ):
58+ """
59+ Enter the context for this runner.
60+ :return: This runner instance.
61+ :rtype: TargetRunner
62+ """
4963 return self
5064
5165 def __exit__ (self , * _ ):
66+ """
67+ Exit the context for this runner.
68+ """
5269 pass
5370
5471
@@ -77,20 +94,20 @@ class SubprocessTargetRunner(TargetRunner):
7794 """
7895
7996 def __init__ (self ,
80- runner_cmd ,
81- target ,
82- connect = True ,
83- boot_timeout = 60 ):
97+ runner_cmd : SubprocessCommand ,
98+ target : Target ,
99+ connect : bool = True ,
100+ boot_timeout : int = 60 ):
84101 super ().__init__ (target = target )
85102
86- self .boot_timeout = boot_timeout
103+ self .boot_timeout : int = boot_timeout
87104
88105 self .logger .info ('runner_cmd: %s' , runner_cmd )
89106
90107 try :
91108 self .runner_process = get_subprocess (runner_cmd )
92109 except Exception as ex :
93- raise HostError (f'Error while running "{ runner_cmd } ": { ex } ' ) from ex
110+ raise HostError (f'Error while running "{ runner_cmd !r } ": { ex } ' ) from ex
94111
95112 if connect :
96113 self .wait_boot_complete ()
@@ -107,16 +124,16 @@ def __exit__(self, *_):
107124
108125 self .terminate ()
109126
110- def wait_boot_complete (self ):
127+ def wait_boot_complete (self ) -> None :
111128 """
112- Wait for target OS to finish boot up and become accessible over SSH in at most
113- ``SubprocessTargetRunner. boot_timeout` ` seconds.
129+ Wait for the target OS to finish booting and become accessible within
130+ :attr:` boot_timeout` seconds.
114131
115- :raises TargetStableError: In case of timeout.
132+ :raises TargetStableError: If the target is inaccessible after the timeout.
116133 """
117134
118135 start_time = time .time ()
119- elapsed = 0
136+ elapsed : float = 0. 0
120137 while self .boot_timeout >= elapsed :
121138 try :
122139 self .target .connect (timeout = self .boot_timeout - elapsed )
@@ -132,9 +149,9 @@ def wait_boot_complete(self):
132149 self .terminate ()
133150 raise TargetStableError (f'Target is inaccessible for { self .boot_timeout } seconds!' )
134151
135- def terminate (self ):
152+ def terminate (self ) -> None :
136153 """
137- Terminate ``SubprocessTargetRunner.runner_process`` .
154+ Terminate the subprocess associated with this runner .
138155 """
139156
140157 self .logger .debug ('Killing target runner...' )
@@ -150,7 +167,7 @@ class NOPTargetRunner(TargetRunner):
150167 :type target: Target
151168 """
152169
153- def __init__ (self , target ) :
170+ def __init__ (self , target : Target ) -> None :
154171 super ().__init__ (target = target )
155172
156173 def __enter__ (self ):
@@ -159,11 +176,63 @@ def __enter__(self):
159176 def __exit__ (self , * _ ):
160177 pass
161178
162- def terminate (self ):
179+ def terminate (self ) -> None :
163180 """
164181 Nothing to terminate for NOP target runners.
165182 Defined to be compliant with other runners (e.g., ``SubprocessTargetRunner``).
166183 """
184+ pass
185+
186+
187+ QEMUTargetUserSettings = TypedDict ("QEMUTargetUserSettings" , {
188+ 'kernel_image' : str ,
189+ 'arch' : NotRequired [str ],
190+ 'cpu_type' : NotRequired [str ],
191+ 'initrd_image' : str ,
192+ 'mem_size' : NotRequired [int ],
193+ 'num_cores' : NotRequired [int ],
194+ 'num_threads' : NotRequired [int ],
195+ 'cmdline' : NotRequired [str ],
196+ 'enable_kvm' : NotRequired [bool ],
197+ })
198+
199+ QEMUTargetRunnerSettings = TypedDict ("QEMUTargetRunnerSettings" , {
200+ 'kernel_image' : str ,
201+ 'arch' : str ,
202+ 'cpu_type' : str ,
203+ 'initrd_image' : str ,
204+ 'mem_size' : int ,
205+ 'num_cores' : int ,
206+ 'num_threads' : int ,
207+ 'cmdline' : str ,
208+ 'enable_kvm' : bool ,
209+ })
210+
211+
212+ SshConnectionSettings = TypedDict ("SshConnectionSettings" , {
213+ 'username' : str ,
214+ 'password' : str ,
215+ 'keyfile' : Optional [Union [LiteralString , StrPath , BytesPath ]],
216+ 'host' : str ,
217+ 'port' : int ,
218+ 'timeout' : float ,
219+ 'platform' : 'Platform' ,
220+ 'sudo_cmd' : str ,
221+ 'strict_host_check' : bool ,
222+ 'use_scp' : bool ,
223+ 'poll_transfers' : bool ,
224+ 'start_transfer_poll_delay' : int ,
225+ 'total_transfer_timeout' : int ,
226+ 'transfer_poll_period' : int ,
227+ })
228+
229+
230+ class QEMUTargetRunnerTargetFactory (Protocol ):
231+ """
232+ Protocol for Lambda function for creating :class:`Target` based object.
233+ """
234+ def __call__ (self , * , connect : bool , conn_cls , connection_settings : SshConnectionSettings ) -> Target :
235+ ...
167236
168237
169238class QEMUTargetRunner (SubprocessTargetRunner ):
@@ -177,7 +246,7 @@ class QEMUTargetRunner(SubprocessTargetRunner):
177246
178247 * ``arch``: Architecture type. Defaults to ``aarch64``.
179248
180- * ``cpu_types ``: List of CPU ids for QEMU. The list only contains ``cortex-a72`` by
249+ * ``cpu_type ``: List of CPU ids for QEMU. The list only contains ``cortex-a72`` by
181250 default. This parameter is valid for Arm architectures only.
182251
183252 * ``initrd_image``: This points to the location of initrd image (e.g.,
@@ -212,21 +281,25 @@ class QEMUTargetRunner(SubprocessTargetRunner):
212281 """
213282
214283 def __init__ (self ,
215- qemu_settings ,
216- connection_settings = None ,
217- make_target = LinuxTarget ,
218- ** args ):
284+ qemu_settings : QEMUTargetUserSettings ,
285+ connection_settings : Optional [ SshUserConnectionSettings ] = None ,
286+ make_target : QEMUTargetRunnerTargetFactory = cast ( QEMUTargetRunnerTargetFactory , LinuxTarget ) ,
287+ ** args ) -> None :
219288
220- self . connection_settings = {
289+ default_connection_settings = {
221290 'host' : '127.0.0.1' ,
222291 'port' : 8022 ,
223292 'username' : 'root' ,
224293 'password' : 'root' ,
225294 'strict_host_check' : False ,
226295 }
227- self .connection_settings = {** self .connection_settings , ** (connection_settings or {})}
228296
229- qemu_args = {
297+ self .connection_settings : SshConnectionSettings = cast (SshConnectionSettings , {
298+ ** default_connection_settings ,
299+ ** (connection_settings or {})
300+ })
301+
302+ qemu_default_args = {
230303 'arch' : 'aarch64' ,
231304 'cpu_type' : 'cortex-a72' ,
232305 'mem_size' : 512 ,
@@ -235,7 +308,7 @@ def __init__(self,
235308 'cmdline' : 'console=ttyAMA0' ,
236309 'enable_kvm' : True ,
237310 }
238- qemu_args = {** qemu_args , ** qemu_settings }
311+ qemu_args : QEMUTargetRunnerSettings = cast ( QEMUTargetRunnerSettings , {** qemu_default_args , ** qemu_settings })
239312
240313 qemu_executable = f'qemu-system-{ qemu_args ["arch" ]} '
241314 qemu_path = which (qemu_executable )
0 commit comments