Difficult-Rocket/Difficult_Rocket/utils/thread.py

#  -------------------------------
#  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")