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
« prev ^ index » next coverage.py v7.4.1, created at 2024-01-31 11:23 +0100
1"""
2Very simple threading abstraction.
3"""
5import threading
6import typing
8from result import Err, Ok, Result
9from typing_extensions import Self
11P = typing.ParamSpec("P")
12R = typing.TypeVar("R")
15class ThreadWithReturn(typing.Generic[R], threading.Thread):
16 """
17 Should not be used directly.
19 Rather use the @thread decorator,
20 which changes the return type of function() -> T into function() -> ThreadWithReturn[T]
21 """
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]]
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.
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 = []
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
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()
64 def result(self) -> "Result[R, Exception | None]":
65 """
66 Get the result value (Ok or Err) from the threaded function.
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)
80 def then(self, callback: typing.Callable[[R], R]) -> Self:
81 """
82 Attach a callback (which runs in the thread as well) on success.
84 Returns 'self' so you can do .then().then().then().
85 """
86 self._callbacks.append(callback)
87 return self # for builder pattern
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.
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)
98 return self
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)
106 match self.result():
107 case Ok(value):
108 return value
109 case Err(exc):
110 raise exc or Exception("Something went wrong.")
112 # thread must be ready so Err(None) can't happen
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 """
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
127 return wraps
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 """
138 def wraps(inner_function: typing.Callable[P, R]) -> typing.Callable[P, ThreadWithReturn[R]]:
139 """Idem ditto."""
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
147 return inner
149 return wraps
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!
161 Examples:
162 @thread
163 def myfunc():
164 ...
166 @thread()
167 def otherfunc():
168 ...
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:
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
184 return wraps
186 else:
187 return thread
190__all__ = [
191 "ThreadWithReturn",
192 "thread",
193]