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

65 statements  

« prev     ^ index     » next       coverage.py v7.4.3, created at 2024-02-29 21:18 +0100

1""" 

2Very simple threading abstraction. 

3""" 

4 

5import functools 

6import threading 

7import typing 

8 

9from result import Err, Ok, Result 

10from typing_extensions import Self 

11 

12P = typing.ParamSpec("P") 

13R = typing.TypeVar("R") 

14 

15 

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

17 """ 

18 Should not be used directly. 

19 

20 Rather use the @thread decorator, 

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

22 """ 

23 

24 _target: typing.Callable[P, R] 

25 _args: P.args 

26 _kwargs: P.kwargs 

27 _return: R | Exception 

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

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

30 

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

32 """ 

33 Setup callbacks, otherwise same logic as super. 

34 

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

36 """ 

37 kw["target"] = target 

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

39 self._callbacks = [] 

40 self._catch = [] 

41 

42 def run(self) -> None: 

43 """ 

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

45 """ 

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

47 return 

48 

49 try: 

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

51 for callback in self._callbacks: 

52 result = callback(result) 

53 self._return = result 

54 except Exception as _e: 

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

56 for err_callback in self._catch: 

57 e = err_callback(e) 

58 self._return = e 

59 finally: 

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

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

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

63 # keep self._return for .result() 

64 

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

66 """ 

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

68 

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

70 """ 

71 if self.is_alive(): 

72 # still busy 

73 return Err(None) 

74 else: 

75 result = self._return 

76 if isinstance(result, Exception): 

77 return Err(result) 

78 else: 

79 return Ok(result) 

80 

81 def is_done(self) -> bool: 

82 """ 

83 Returns whether the thread has finished (result or error). 

84 """ 

85 return not self.is_alive() 

86 

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

88 """ 

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

90 

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

92 """ 

93 self._callbacks.append(callback) 

94 return self # for builder pattern 

95 

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

97 """ 

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

99 

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

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

102 """ 

103 self._catch.append(callback) 

104 

105 return self 

106 

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

108 """ 

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

110 """ 

111 super().join(timeout) 

112 

113 match self.result(): 

114 case Ok(value): 

115 return value 

116 case Err(exc): 

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

118 

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

120 

121 

122@typing.overload 

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

124 """ 

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

126 """ 

127 

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

129 """Idem ditto.""" 

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

131 my_thread.start() 

132 return my_thread 

133 

134 return wraps 

135 

136 

137@typing.overload 

138def thread( 

139 my_function: None = None, 

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

141 """ 

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

143 """ 

144 

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

146 """Idem ditto.""" 

147 

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

149 """Idem ditto.""" 

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

151 my_thread.start() 

152 return my_thread 

153 

154 return inner 

155 

156 return wraps 

157 

158 

159def thread( 

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

161) -> ( 

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

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

164): 

165 """ 

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

167 

168 Examples: 

169 @thread 

170 def myfunc(): 

171 ... 

172 

173 @thread() 

174 def otherfunc(): 

175 ... 

176 

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

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

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

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

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

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

183 """ 

184 if my_function is None: 

185 return thread 

186 

187 @functools.wraps(my_function) 

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

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

190 my_thread.start() 

191 return my_thread 

192 

193 return wraps 

194 

195 

196__all__ = [ 

197 "ThreadWithReturn", 

198 "thread", 

199]