funcs.py

from __future__ import division, print_function, absolute_import
from collections.abc import Mapping, Iterable
from functools import wraps
import traceback
import sys
import gc
import signal
import inspect
import types
from concurrent.futures import ThreadPoolExecutor, TimeoutError

__all__ = [
    "unpack_args",
    "callback_gc",
    "catch_exception",
    "full_traceback",
    "timeout",
    "hotpatch_module",
    "print_flush",
]


def get_default_args(func):
    """
    cf. https://stackoverflow.com/a/12627202/
    """
    signature = inspect.signature(func)
    return {
        k: v.default
        for k, v in signature.parameters.items()
        if v.default is not inspect.Parameter.empty
    }


def unpack_args(func):
    @wraps(func)
    def wrapper(args):
        if isinstance(args, Mapping):
            return func(**args)
        elif isinstance(args, Iterable):
            return func(*args)
        else:
            return func(args)

    return wrapper


def callback_gc(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        res = func(*args, **kwargs)
        gc.collect()
        return res

    return wrapper


def catch_exception(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        try:
            return func(*args, **kwargs)
        except Exception as e:
            msg = "Failed:\n  {}(*{}, **{})\n{}".format(
                func.__name__, args, kwargs, traceback.format_exc()
            )
            print(msg)
            return e

    return wrapper


def full_traceback(func):
    """
    Seems to not not necessary in Python 3
    http://stackoverflow.com/a/29442282
    http://bugs.python.org/issue13831
    """

    @wraps(func)
    def wrapper(*args, **kwargs):
        try:
            return func(*args, **kwargs)
        except Exception as e:
            msg = "{}\n\nOriginal {}".format(e, traceback.format_exc())
            raise type(e)(msg)

    return wrapper


class timeout:
    """
    Handling timeout.

    Note that `signal.signal` can only be called from the main thread
    in Unix-like system, and can not be nested,
    i.e. when timeout is called under another timeout, the first one will overridden.

    To avoid above issues, one should use the decorator mode with `thread=True`.

    Examples
    --------
        import time

        # context mode
        with timeout(seconds=1):
            time.sleep(4)

        # decorator mode
        @timeout(1)
        def func():
            time.sleep(4)
        func()

    Reference
    ---------
    https://stackoverflow.com/a/22348885/ for decorator
    https://stackoverflow.com/a/2282656/ for context
    https://stackoverflow.com/a/11901541/ for `signal.setitimer`
    """

    def __init__(self, seconds=1, exception=TimeoutError("Timeout."), thread=True):
        """
        seconds :
            Note `signal.alarm`
        thread :
            If True, `concurrent.futures.ThreadPoolExecutor` is used,
            otherwise `signal.signal` is used.
            Only takes effect in decorator mode.
        """
        self.seconds = seconds
        self.thread = thread

        if isinstance(exception, type) and issubclass(exception, Exception):
            self.exception = exception("Timeout.")
        elif isinstance(exception, Exception):
            self.exception = exception
        else:
            self.exception = TimeoutError(exception)

    def handler(self, signum, frame):
        raise self.exception

    def __enter__(self):
        signal.signal(signal.SIGALRM, self.handler)
        # signal.alarm(self.seconds)
        signal.setitimer(signal.ITIMER_REAL, self.seconds)

    def __exit__(self, type, value, traceback):
        signal.alarm(0)

    def __call__(self, func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            if self.thread:
                with ThreadPoolExecutor(1) as pool:
                    res = pool.submit(func, *args, **kwargs)
                    res.set_exception(self.exception)
                    return res.result(self.seconds)
            else:
                with self:
                    return func(*args, **kwargs)

        return wrapper


def hotpatch_module(module):
    """
    Return a decorator that rebuilds a function object so that its global
    namespace points to `module.__dict__`, then installs the new function
    into the specified module under the same name.
    """

    def decorator(func):
        # Reconstruct a new function using the original code object and defaults,
        # but with globals set to the target module's namespace.
        new_func = types.FunctionType(
            func.__code__,  # the compiled bytecode of the original function
            module.__dict__,  # use the module's globals for name resolution
            name=func.__name__,  # preserve the original function name
            argdefs=func.__defaults__,  # positional/default argument values
            closure=func.__closure__,  # any closure cells, if the function uses them
        )
        # Copy over keyword-only default arguments, if any
        new_func.__kwdefaults__ = getattr(func, "__kwdefaults__", None)

        # Replace or add the function in the target module under its name
        setattr(module, func.__name__, new_func)
        return new_func

    return decorator


def print_flush(*args, **kwargs):
    flush = kwargs.pop("flush", True)
    print(*args, **kwargs)
    file = kwargs.get("file", sys.stdout)
    if flush and file is not None:
        file.flush()