from __future__ import print_function
import os
import sys
import functools
import contextlib
import hashlib
import numpy
import fcntl

def make_list(generator_func):
    """Create a list from a generator function.
    """
    def wrapper(*args, **kwargs):
        return list(generator_func(*args, **kwargs))
    return functools.update_wrapper(wrapper, generator_func)

def vectorize(generator_func):
    """Create a numpy array from an generator function.
    
    numpy.fromiter could have been used, but then the dtype would
    have to be specified (or maybe guessed based on the first item).
    This version is not super efficient, but easy to use.

    >>> @vectorize
    ... def func():
    ...     yield 1
    ...     yield 2
    ...     yield 3.4
    >>> func()
    array([ 1. ,  2. ,  3.4])
    """
    def wrapper(*args, **kwargs):
        return numpy.array(list(generator_func(*args, **kwargs)))
    return functools.update_wrapper(wrapper, generator_func)

@contextlib.contextmanager
def flocked(file, exclusive=True):
    fcntl.flock(file, fcntl.LOCK_EX if exclusive else fcntl.LOCK_SH)
    try:
        yield
    finally:
        fcntl.flock(file, fcntl.LOCK_UN)

@contextlib.contextmanager
def capture_stdout(file=None):
    """Capture output temporarily.

    Use as:
    >>> import io
    >>> with capture_stdout(io.StringIO()) as c:
    ...    print(u'hello, world')
    >>> print(c.getvalue(), end='')
    hello, world

    Or simply as:
    >>> def invoke_crazy_function_which_spews_garbage_on_stdout():
    ...    print('x' * 1000)
    >>> with capture_stdout():
    ...    invoke_crazy_function_which_spews_garbage_on_stdout()
    """
    stdout = sys.stdout
    if file is None:
        file = open(os.devnull, 'w')
    sys.stdout = file
    try:
        yield file
    finally:
        sys.stdout = stdout

@contextlib.contextmanager
def chdired(dirname):
    """Change directory temporarily"""
    old = os.getcwd()
    os.chdir(dirname)
    try:
        yield
    finally:
        os.chdir(old)

@contextlib.contextmanager
def numpy_seterr(all=None, divide=None, over=None, under=None, invalid=None):
    old = numpy.seterr(all=all, divide=divide, over=over, under=under,
                       invalid=invalid)
    try:
        yield
    finally:
        numpy.seterr(**old)

def disk_cached(basepath):
    """
    Caches results in numpy files on disk.

    Specified basepath must be under the control of the caller,
    and entries must be cleaned manually. So this a basis for
    a solution, not the final thing certainly.

    >>> @disk_cached('/tmp/f_result_')
    ... def f(*args, **kwargs):
    ...     print('f')
    ...     ans = numpy.array([len(args), len(kwargs)])
    ...     return ans
    >>> f(3)
    f
    array([1, 0])
    >>> f(3)
    array([1, 0])
    >>> f(*[3])
    array([1, 0])
    >>> f(a=1, b=2)
    f
    array([0, 2])
    >>> f(b=2, a=1)
    array([0, 2])
    >>> f([1,2,3])
    f
    array([1, 0])
    >>> f([1,2,3])
    array([1, 0])
    """
    def decorator(func):
        def wrapper(*args, **kwargs):
            key = (args, tuple(sorted(kwargs.items())))
            tag = hashlib.sha1(repr(key)).hexdigest()
            path = basepath + tag + '.npy'
            try:
                ans = numpy.load(path)
            except TypeError:
                print('warning: function arguments are unhashable')
                return func(*args, **kwargs)
            except IOError:
                # value is not present in cache
                ans = func(*args, **kwargs)
                numpy.save(path, ans)
            return ans
        return functools.update_wrapper(wrapper, func)
    return decorator