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
« prev ^ index » next coverage.py v7.4.3, created at 2024-02-29 21:18 +0100
1"""
2Very simple threading abstraction.
3"""
5import functools
6import threading
7import typing
9from result import Err, Ok, Result
10from typing_extensions import Self
12P = typing.ParamSpec("P")
13R = typing.TypeVar("R")
16class ThreadWithReturn(typing.Generic[R], threading.Thread):
17 """
18 Should not be used directly.
20 Rather use the @thread decorator,
21 which changes the return type of function() -> T into function() -> ThreadWithReturn[T]
22 """
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]]
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.
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 = []
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
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()
65 def result(self) -> "Result[R, Exception | None]":
66 """
67 Get the result value (Ok or Err) from the threaded function.
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)
81 def is_done(self) -> bool:
82 """
83 Returns whether the thread has finished (result or error).
84 """
85 return not self.is_alive()
87 def then(self, callback: typing.Callable[[R], R]) -> Self:
88 """
89 Attach a callback (which runs in the thread as well) on success.
91 Returns 'self' so you can do .then().then().then().
92 """
93 self._callbacks.append(callback)
94 return self # for builder pattern
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.
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)
105 return self
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)
113 match self.result():
114 case Ok(value):
115 return value
116 case Err(exc):
117 raise exc or Exception("Something went wrong.")
119 # thread must be ready so Err(None) can't happen
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 """
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
134 return wraps
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 """
145 def wraps(inner_function: typing.Callable[P, R]) -> typing.Callable[P, ThreadWithReturn[R]]:
146 """Idem ditto."""
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
154 return inner
156 return wraps
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!
168 Examples:
169 @thread
170 def myfunc():
171 ...
173 @thread()
174 def otherfunc():
175 ...
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
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
193 return wraps
196__all__ = [
197 "ThreadWithReturn",
198 "thread",
199]