Coverage for src/threadful/__init__.py: 100%

61 statements  

« prev     ^ index     » next       coverage.py v7.4.1, created at 2024-01-31 11:23 +0100

1""" 

2Very simple threading abstraction. 

3""" 

4 

5import threading 

6import typing 

7 

8from result import Err, Ok, Result 

9from typing_extensions import Self 

10 

11P = typing.ParamSpec("P") 

12R = typing.TypeVar("R") 

13 

14 

15class ThreadWithReturn(typing.Generic[R], threading.Thread): 

16 """ 

17 Should not be used directly. 

18 

19 Rather use the @thread decorator, 

20 which changes the return type of function() -> T into function() -> ThreadWithReturn[T] 

21 """ 

22 

23 _target: typing.Callable[P, R] 

24 _args: P.args 

25 _kwargs: P.kwargs 

26 _return: R | Exception 

27 _callbacks: list[typing.Callable[[R], R]] 

28 _catch: list[typing.Callable[[Exception | R], Exception | R]] 

29 

30 def __init__(self, target: typing.Callable[P, R], *a: typing.Any, **kw: typing.Any) -> None: 

31 """ 

32 Setup callbacks, otherwise same logic as super. 

33 

34 'target' is explicitly mentioned outside of kw for type hinting. 

35 """ 

36 kw["target"] = target 

37 super().__init__(*a, **kw) 

38 self._callbacks = [] 

39 self._catch = [] 

40 

41 def run(self) -> None: 

42 """ 

43 Called in a new thread and handles the calling logic. 

44 """ 

45 if self._target is None: # pragma: no cover 

46 return 

47 

48 try: 

49 result = self._target(*self._args, **self._kwargs) 

50 for callback in self._callbacks: 

51 result = callback(result) 

52 self._return = result 

53 except Exception as _e: 

54 e: Exception | R = _e # make mypy happy 

55 for err_callback in self._catch: 

56 e = err_callback(e) 

57 self._return = e 

58 finally: 

59 # Avoid a refcycle if the thread is running a function with 

60 # an argument that has a member that points to the thread. 

61 del self._target, self._args, self._kwargs, self._callbacks 

62 # keep self._return for .result() 

63 

64 def result(self) -> "Result[R, Exception | None]": 

65 """ 

66 Get the result value (Ok or Err) from the threaded function. 

67 

68 If the thread is not ready, Err(None) is returned. 

69 """ 

70 if self.is_alive(): 

71 # still busy 

72 return Err(None) 

73 else: 

74 result = self._return 

75 if isinstance(result, Exception): 

76 return Err(result) 

77 else: 

78 return Ok(result) 

79 

80 def then(self, callback: typing.Callable[[R], R]) -> Self: 

81 """ 

82 Attach a callback (which runs in the thread as well) on success. 

83 

84 Returns 'self' so you can do .then().then().then(). 

85 """ 

86 self._callbacks.append(callback) 

87 return self # for builder pattern 

88 

89 def catch(self, callback: typing.Callable[[Exception | R], Exception | R]) -> Self: 

90 """ 

91 Attach a callback (which runs in the thread as well) on error. 

92 

93 You can either return a new Exception or a fallback value. 

94 Returns 'self' so you can do .then().then().then(). 

95 """ 

96 self._catch.append(callback) 

97 

98 return self 

99 

100 def join(self, timeout: int | float | None = None) -> R: # type: ignore 

101 """ 

102 Enhanced version of thread.join that also returns the value or raises the exception. 

103 """ 

104 super().join(timeout) 

105 

106 match self.result(): 

107 case Ok(value): 

108 return value 

109 case Err(exc): 

110 raise exc or Exception("Something went wrong.") 

111 

112 # thread must be ready so Err(None) can't happen 

113 

114 

115@typing.overload 

116def thread(my_function: typing.Callable[P, R]) -> typing.Callable[P, ThreadWithReturn[R]]: # pragma: no cover 

117 """ 

118 Code in this function is never executed, just shown for reference of the complex return type. 

119 """ 

120 

121 def wraps(*a: P.args, **kw: P.kwargs) -> ThreadWithReturn[R]: 

122 """Idem ditto.""" 

123 my_thread = ThreadWithReturn(target=my_function, args=a, kwargs=kw) 

124 my_thread.start() 

125 return my_thread 

126 

127 return wraps 

128 

129 

130@typing.overload 

131def thread( 

132 my_function: None = None, 

133) -> typing.Callable[[typing.Callable[P, R]], typing.Callable[P, ThreadWithReturn[R]]]: # pragma: no cover 

134 """ 

135 Code in this function is never executed, just shown for reference of the complex return type. 

136 """ 

137 

138 def wraps(inner_function: typing.Callable[P, R]) -> typing.Callable[P, ThreadWithReturn[R]]: 

139 """Idem ditto.""" 

140 

141 def inner(*a: P.args, **kw: P.kwargs) -> ThreadWithReturn[R]: 

142 """Idem ditto.""" 

143 my_thread = ThreadWithReturn(target=inner_function, args=a, kwargs=kw) 

144 my_thread.start() 

145 return my_thread 

146 

147 return inner 

148 

149 return wraps 

150 

151 

152def thread( 

153 my_function: typing.Callable[P, R] | None = None 

154) -> ( 

155 typing.Callable[[typing.Callable[P, R]], typing.Callable[P, ThreadWithReturn[R]]] 

156 | typing.Callable[P, ThreadWithReturn[R]] 

157): 

158 """ 

159 This decorator can be used to automagically make functions threaded! 

160 

161 Examples: 

162 @thread 

163 def myfunc(): 

164 ... 

165 

166 @thread() 

167 def otherfunc(): 

168 ... 

169 

170 myfunc() and otherfunc() now return a custom thread object, 

171 from which you can get the result value or exception with .result(). 

172 This uses a Result (Ok or Err) type from rustedpy/result (based on the Rust Result type.) 

173 If the thread is not done yet, it will return Err(None) 

174 You can also call .join(), which waits (blocking) until the thread is done 

175 and then returns the return value or raises an exception (if raised in the thread) 

176 """ 

177 if my_function is not None: 

178 

179 def wraps(*a: P.args, **kw: P.kwargs) -> ThreadWithReturn[R]: 

180 my_thread = ThreadWithReturn(target=my_function, args=a, kwargs=kw) 

181 my_thread.start() 

182 return my_thread 

183 

184 return wraps 

185 

186 else: 

187 return thread 

188 

189 

190__all__ = [ 

191 "ThreadWithReturn", 

192 "thread", 

193]