def wrapped(
arg: typing.Union[selectable.LabHost, linux.LinuxShell, None] = None,
*args: typing.Any,
**kwargs: typing.Any,
) -> typing.Any:
with contextlib.ExitStack() as cx:
lh: selectable.LabHost
lnx: linux.LinuxShell
# Acquire LabHost
if arg is None:
lh = cx.enter_context(selectable.acquire_lab())
elif isinstance(arg, linux.Lab):
lh = cx.enter_context(arg) # type: ignore
elif not isinstance(arg, linux.LinuxShell):
raise TypeError(
f"Argument to {tc!r} must either be a lab-host or a board linux (found {arg!r})"
)
# Acquire Linux
if arg is None or isinstance(arg, linux.Lab):
b = cx.enter_context(selectable.acquire_board(lh))
lnx = cx.enter_context(selectable.acquire_linux(b))
else:
lnx = cx.enter_context(arg)
return tc(lnx, *args, **kwargs)
def wrapped(
arg: typing.Union[selectable.LabHost, board.UBootShell, None] = None,
*args: typing.Any,
**kwargs: typing.Any,
) -> typing.Any:
with contextlib.ExitStack() as cx:
lh: selectable.LabHost
ub: board.UBootShell
# Acquire LabHost
if arg is None:
lh = cx.enter_context(selectable.acquire_lab())
elif isinstance(arg, linux.Lab):
lh = cx.enter_context(arg)
elif not isinstance(arg, board.UBootShell):
raise TypeError(
f"Argument to {tc!r} must either be a lab-host or a UBootShell (found {arg!r})"
)
# Acquire U-Boot
if isinstance(arg, board.UBootShell):
ub = cx.enter_context(arg)
else:
b = cx.enter_context(selectable.acquire_board(lh))
ub = cx.enter_context(selectable.acquire_uboot(b))
return tc(ub, *args, **kwargs)