DivisionBy-Zero/erpa-sweng

import logging
import shutil
import tempfile
from contextlib import contextmanager
from functools import wraps
from inspect import isgeneratorfunction, unwrap
from typing import Any, Callable, cast, Dict, Generator, IO, Tuple, TypeVar

G = TypeVar('G')
S = TypeVar('S')


def is_context_manager(f):
    return isgeneratorfunction(unwrap(getattr(f, '__wrapped__', None), stop=isgeneratorfunction))


CTX = Any  # Python3.5's typing does not known about type ContextManager[G]
CTX_GENERATOR = Callable[..., CTX]
CTX_GEN_ARGS = Tuple[Any, ...]
CTX_GEN_KWARGS = Dict[str, Any]
DECORATOR_FN = Callable[[Callable[..., S]], Callable[..., S]]
FN_ARGS = Tuple[Any, ...]
FN_KWARGS = Dict[str, Any]


# Meaningful argument names are helpful: pylint: disable=invalid-name
def mk_with_context_decorator(
        ctx_generator: CTX_GENERATOR,
        ctx_generator_args_from_fn_args: Callable[[FN_ARGS, FN_KWARGS], Tuple[CTX_GEN_ARGS, CTX_GEN_KWARGS]],
        fn_args_from_fn_args_and_ctx: Callable[[FN_ARGS, FN_KWARGS, CTX], Tuple[FN_ARGS, FN_KWARGS]]) -> DECORATOR_FN:
    """Make a decorator that wraps a function or contextmanager within the context provided by `ctx_generator`."""
    def decorator(f: Callable[..., S]) -> Callable[..., S]:
        def function_wrapper(*args, **kwargs) -> S:
            context_generator_args, context_generator_kwargs = ctx_generator_args_from_fn_args(args, kwargs)
            with ctx_generator(*context_generator_args, **context_generator_kwargs) as ct:
                args, kwargs = fn_args_from_fn_args_and_ctx(args, kwargs, ct)
                return f(*args, **kwargs)

        @contextmanager
        def generator_wrapper(*args, **kwargs) -> Generator[G, None, None]:
            context_generator_args, context_generator_kwargs = ctx_generator_args_from_fn_args(args, kwargs)
            with ctx_generator(*context_generator_args, **context_generator_kwargs) as ct:
                args, kwargs = fn_args_from_fn_args_and_ctx(args, kwargs, ct)
                with cast(CTX, f(*args, **kwargs)) as result:
                    yield result
        return wraps(f)(generator_wrapper if is_context_manager(f) else function_wrapper)
    return decorator


def with_context(context_generator: CTX_GENERATOR, *context_generator_args,
                 **context_generator_kwargs) -> DECORATOR_FN:
    """Returns a decorator that calls the function within a context. The context is passed as last unnamed argument.
    The context is generated by calling context_generator with its unnamed and named arguments.
    """
    def ctx_generator_args_from_fn_args(*_):
        return context_generator_args, context_generator_kwargs

    def fn_args_from_fn_args_and_ctx(fn_args, fn_kwargs, ctx):
        return fn_args + cast(tuple, ctx if isinstance(ctx, tuple) else (ctx,)), fn_kwargs
    return mk_with_context_decorator(context_generator, ctx_generator_args_from_fn_args, fn_args_from_fn_args_and_ctx)


def with_context_using_instance(context_generator: CTX_GENERATOR, *context_generator_args,
                                **context_generator_kwargs) -> DECORATOR_FN:
    """Returns a decorator that calls the function within a context. The context is passed as last unnamed argument.
    The context is generated by calling context_generator with function's first unnamed argument (ie. self),
    along with its unnamed and named arguments.
    """
    def ctx_generator_args_from_fn_args(fn_args, _):
        return (fn_args[0],) + context_generator_args, context_generator_kwargs

    def fn_args_from_fn_args_and_ctx(fn_args, fn_kwargs, ctx):
        return fn_args + cast(tuple, ctx if isinstance(ctx, tuple) else (ctx,)), fn_kwargs
    return mk_with_context_decorator(context_generator, ctx_generator_args_from_fn_args, fn_args_from_fn_args_and_ctx)


@contextmanager
def tmp_directory_ctx_generator(maybe_self=None, prefix='erpa-tmp-') -> Generator[str, None, None]:
    """Creates and yields a temporary directory which is destroyed when the contexts quits.
    If maybe_self is specified and has property `tmpdir`, it will be use as basedir.

    See also
    :func:`tempfile.NamedTemporaryFile`
    """
    working_dir = getattr(maybe_self, 'tmpdir', None)
    tmpdir = tempfile.mkdtemp(dir=working_dir, prefix=prefix)
    try:
        yield tmpdir
    finally:
        try:
            shutil.rmtree(tmpdir)
        except Exception as exception:
            logging.warning('Exception while deleting tmpdir {}. Retrying. Exception: {}'.format(tmpdir, exception))
            try:
                shutil.rmtree(tmpdir, ignore_errors=True)
            except Exception as exc:
                logging.error('Leaving orphan tree: Could not forcibly delete tmpdir {}: {}'.format(tmpdir, exc))


@contextmanager
def tmp_file_ctx_generator(maybe_self=None, prefix='erpa-tmp-') -> Generator[IO, None, None]:
    """Creates and yields a temporary file which is destroyed when the contexts quits.
    If maybe_self is specified and has property `tmpdir`, it will be use as basedir.

    See also
    :func:`tempfile.NamedTemporaryFile`
    """
    working_dir = getattr(maybe_self, 'tmpdir', None)
    with tempfile.NamedTemporaryFile(buffering=False, dir=working_dir, prefix=prefix) as tmpfile:
        yield tmpfile