206 lines
6.5 KiB
Python
206 lines
6.5 KiB
Python
# -------------------------------
|
||
# Difficult Rocket
|
||
# Copyright © 2020-2023 by shenjackyuanjie 3695888@qq.com
|
||
# All rights reserved
|
||
# -------------------------------
|
||
|
||
"""
|
||
writen by shenjackyuanjie
|
||
mail: 3695888@qq.com
|
||
github: @shenjackyuanjie
|
||
gitee: @shenjackyuanjie
|
||
"""
|
||
|
||
import functools
|
||
import inspect
|
||
import threading
|
||
from threading import Lock
|
||
from typing import Optional, Callable, Union, List
|
||
|
||
from Difficult_Rocket.exception.threading import LockTimeOutError
|
||
|
||
__all__ = [
|
||
"new_thread",
|
||
"FunctionThread",
|
||
"ThreadLock",
|
||
"record_thread",
|
||
"record_destination",
|
||
]
|
||
|
||
record_thread = False
|
||
record_destination: List[Callable[["FunctionThread"], None]] = []
|
||
|
||
|
||
class ThreadLock:
|
||
def __init__(self, the_lock: Lock, time_out: Union[float, int] = 1 / 60) -> None:
|
||
self.lock = the_lock
|
||
self.time_out = time_out
|
||
|
||
def __enter__(self):
|
||
self.lock.acquire(timeout=self.time_out)
|
||
if not self.lock.locked():
|
||
raise LockTimeOutError(f"Lock time Out with {self.time_out}")
|
||
return self
|
||
|
||
def __exit__(self, exc_type, exc_val, exc_tb):
|
||
if self.lock.locked():
|
||
self.lock.release()
|
||
|
||
|
||
"""
|
||
This part of code come from MCDReforged(https://github.com/Fallen-Breath/MCDReforged)
|
||
Very thanks to Fallen_Breath and other coder who helped MCDR worked better
|
||
GNU Lesser General Public License v3.0（GNU LGPL v3)
|
||
(have some changes)
|
||
"""
|
||
|
||
|
||
def copy_signature(target: Callable, origin: Callable) -> Callable:
|
||
"""
|
||
Copy the function signature of origin into target
|
||
"""
|
||
# https://stackoverflow.com/questions/39926567/python-create-decorator-preserving-function-arguments
|
||
target.__signature__ = inspect.signature(origin)
|
||
return target
|
||
|
||
|
||
class FunctionThread(threading.Thread):
|
||
"""
|
||
A Thread subclass which is used in decorator :func:`new_thread` to wrap a synchronized function call
|
||
"""
|
||
|
||
__NONE = object()
|
||
|
||
def __init__(self, target, name, args, kwargs, daemon):
|
||
super().__init__(
|
||
target=target, args=args, kwargs=kwargs, name=name, daemon=daemon
|
||
)
|
||
self.__return_value = self.__NONE
|
||
self.__error = None
|
||
|
||
def wrapped_target(*args_, **kwargs_):
|
||
try:
|
||
self.__return_value = target(*args_, **kwargs_)
|
||
except Exception as e:
|
||
self.__error = e
|
||
raise e from None
|
||
|
||
self._target = wrapped_target
|
||
|
||
def get_return_value(self, block: bool = False, timeout: Optional[float] = None):
|
||
"""
|
||
Get the return value of the original function
|
||
|
||
If an exception has occurred during the original function call, the exception will be risen again here
|
||
|
||
Examples::
|
||
|
||
>>> import time
|
||
>>> @new_thread
|
||
... def do_something(text: str):
|
||
... time.sleep(1)
|
||
... return text
|
||
|
||
>>> do_something('task').get_return_value(block=True)
|
||
'task'
|
||
|
||
:param block: If it should join the thread before getting the return value to make sure the function invocation finishes
|
||
:param timeout: The maximum timeout for the thread join
|
||
:raise RuntimeError: If the thread is still alive when getting return value. Might be caused by ``block=False``
|
||
while the thread is still running, or thread join operation times out
|
||
:return: The return value of the original function
|
||
"""
|
||
if block:
|
||
self.join(timeout)
|
||
if self.__return_value is self.__NONE:
|
||
if self.is_alive():
|
||
raise RuntimeError("The thread is still running")
|
||
raise self.__error
|
||
return self.__return_value
|
||
|
||
|
||
def new_thread(
|
||
arg: Optional[Union[str, Callable]] = None,
|
||
daemon: bool = False,
|
||
log_thread: bool = True,
|
||
):
|
||
"""
|
||
This is a one line solution to make your function executes in parallels.
|
||
When decorated with this decorator, functions will be executed in a new daemon thread
|
||
|
||
This decorator only changes the return value of the function to the created ``Thread`` object.
|
||
Beside the return value, it reserves all signatures of the decorated function,
|
||
so you can safely use the decorated function as if there's no decorating at all
|
||
|
||
It's also a simple compatible upgrade method for old MCDR 0.x plugins
|
||
|
||
The return value of the decorated function is changed to the ``Thread`` object that executes this function
|
||
|
||
The decorated function has 1 extra field:
|
||
|
||
* ``original`` field: The original undecorated function
|
||
|
||
Examples::
|
||
|
||
>>> import time
|
||
|
||
>>> @new_thread('My Plugin Thread')
|
||
... def do_something(text: str):
|
||
... time.sleep(1)
|
||
... print(threading.current_thread().name)
|
||
>>> callable(do_something.original)
|
||
True
|
||
>>> t = do_something('foo')
|
||
>>> isinstance(t, FunctionThread)
|
||
True
|
||
>>> t.join()
|
||
My Plugin Thread
|
||
|
||
:param arg: A :class:`str`, the name of the thread. It's recommend to specify the thread name, so when you
|
||
log something by ``server.logger``, a meaningful thread name will be displayed
|
||
instead of a plain and meaningless ``Thread-3``
|
||
:param daemon: If the thread should be a daemon thread
|
||
:param log_thread: If the thread should be logged to callback defined in record_destination
|
||
"""
|
||
|
||
def wrapper(func):
|
||
@functools.wraps(func) # to preserve the origin function information
|
||
def wrap(*args, **kwargs):
|
||
thread = FunctionThread(
|
||
target=func, args=args, kwargs=kwargs, name=thread_name, daemon=daemon
|
||
)
|
||
if record_thread:
|
||
for destination in record_destination:
|
||
destination(thread)
|
||
thread.start()
|
||
return thread
|
||
|
||
# bring the signature of the func to the wrap function
|
||
# so inspect.getfullargspec(func) works correctly
|
||
copy_signature(wrap, func)
|
||
wrap.original = func # access this field to get the original function
|
||
return wrap
|
||
|
||
# Directly use @new_thread without ending brackets case, e.g. @new_thread
|
||
if isinstance(arg, Callable):
|
||
thread_name = None
|
||
return wrapper(arg)
|
||
# Use @new_thread with ending brackets case, e.g. @new_thread('A'), @new_thread()
|
||
else:
|
||
thread_name = arg
|
||
return wrapper
|
||
|
||
|
||
if __name__ == "__main__":
|
||
from Difficult_Rocket.exception import TestError
|
||
|
||
thread_lock = Lock()
|
||
|
||
test_lock = ThreadLock(thread_lock)
|
||
with test_lock:
|
||
print("do some thing")
|
||
...
|
||
with test_lock:
|
||
print("do some error")
|
||
raise TestError("ah lock test")
|