Coverage for /Users/ajo/work/jumpstarter/jumpstarter/packages/jumpstarter/jumpstarter/client/base.py: 69%
36 statements
« prev ^ index » next coverage.py v7.8.0, created at 2025-05-06 10:21 +0200
« prev ^ index » next coverage.py v7.8.0, created at 2025-05-06 10:21 +0200
1"""
2Base classes for drivers and driver clients
3"""
5from __future__ import annotations
7from contextlib import ExitStack, contextmanager
8from dataclasses import field
10from anyio.from_thread import BlockingPortal
11from pydantic import ConfigDict
12from pydantic.dataclasses import dataclass
14from .core import AsyncDriverClient
15from jumpstarter.streams.blocking import BlockingStream
18@dataclass(kw_only=True, config=ConfigDict(arbitrary_types_allowed=True))
19class DriverClient(AsyncDriverClient):
20 """Base class for driver clients
22 Client methods can be implemented as regular functions,
23 and call the `call` or `streamingcall` helpers internally
24 to invoke exported methods on the driver.
26 Additional client functionalities such as raw stream
27 connections or sharing client-side resources can be added
28 by inheriting mixin classes under `jumpstarter.drivers.mixins`
29 """
31 children: dict[str, DriverClient] = field(default_factory=dict)
33 portal: BlockingPortal
34 stack: ExitStack
36 def call(self, method, *args):
37 """
38 Invoke driver call
40 :param str method: method name of driver call
41 :param list[Any] args: arguments for driver call
43 :return: driver call result
44 :rtype: Any
45 """
46 return self.portal.call(self.call_async, method, *args)
48 def streamingcall(self, method, *args):
49 """
50 Invoke streaming driver call
52 :param str method: method name of streaming driver call
53 :param list[Any] args: arguments for streaming driver call
55 :return: streaming driver call result
56 :rtype: Generator[Any, None, None]
57 """
58 generator = self.portal.call(self.streamingcall_async, method, *args)
59 while True:
60 try:
61 yield self.portal.call(generator.__anext__)
62 except StopAsyncIteration:
63 break
65 @contextmanager
66 def stream(self, method="connect"):
67 """
68 Open a blocking stream session with a context manager.
70 :param str method: method name of streaming driver call
72 :return: blocking stream session object context manager.
73 """
75 with self.portal.wrap_async_context_manager(self.stream_async(method)) as stream:
76 yield BlockingStream(stream=stream, portal=self.portal)
78 @contextmanager
79 def log_stream(self):
80 with self.portal.wrap_async_context_manager(self.log_stream_async()):
81 yield
83 def open_stream(self) -> BlockingStream:
84 """
85 Open a blocking stream session without a context manager.
87 :return: blocking stream session object.
88 :rtype: BlockingStream
89 """
90 return self.stack.enter_context(self.stream())
92 def close(self):
93 """
94 Close the open stream session without a context manager.
95 """
96 self.stack.close()
98 def __del__(self):
99 self.close()