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

1""" 

2Base classes for drivers and driver clients 

3""" 

4 

5from __future__ import annotations 

6 

7from contextlib import ExitStack, contextmanager 

8from dataclasses import field 

9 

10from anyio.from_thread import BlockingPortal 

11from pydantic import ConfigDict 

12from pydantic.dataclasses import dataclass 

13 

14from .core import AsyncDriverClient 

15from jumpstarter.streams.blocking import BlockingStream 

16 

17 

18@dataclass(kw_only=True, config=ConfigDict(arbitrary_types_allowed=True)) 

19class DriverClient(AsyncDriverClient): 

20 """Base class for driver clients 

21 

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. 

25 

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 """ 

30 

31 children: dict[str, DriverClient] = field(default_factory=dict) 

32 

33 portal: BlockingPortal 

34 stack: ExitStack 

35 

36 def call(self, method, *args): 

37 """ 

38 Invoke driver call 

39 

40 :param str method: method name of driver call 

41 :param list[Any] args: arguments for driver call 

42 

43 :return: driver call result 

44 :rtype: Any 

45 """ 

46 return self.portal.call(self.call_async, method, *args) 

47 

48 def streamingcall(self, method, *args): 

49 """ 

50 Invoke streaming driver call 

51 

52 :param str method: method name of streaming driver call 

53 :param list[Any] args: arguments for streaming driver call 

54 

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 

64 

65 @contextmanager 

66 def stream(self, method="connect"): 

67 """ 

68 Open a blocking stream session with a context manager. 

69 

70 :param str method: method name of streaming driver call 

71 

72 :return: blocking stream session object context manager. 

73 """ 

74 

75 with self.portal.wrap_async_context_manager(self.stream_async(method)) as stream: 

76 yield BlockingStream(stream=stream, portal=self.portal) 

77 

78 @contextmanager 

79 def log_stream(self): 

80 with self.portal.wrap_async_context_manager(self.log_stream_async()): 

81 yield 

82 

83 def open_stream(self) -> BlockingStream: 

84 """ 

85 Open a blocking stream session without a context manager. 

86 

87 :return: blocking stream session object. 

88 :rtype: BlockingStream 

89 """ 

90 return self.stack.enter_context(self.stream()) 

91 

92 def close(self): 

93 """ 

94 Close the open stream session without a context manager. 

95 """ 

96 self.stack.close() 

97 

98 def __del__(self): 

99 self.close()