from sympy.utilities.decorator import threaded, xthreaded

from sympy import symbols, Eq, Matrix

from sympy.abc import x, y
from sympy.core.decorators import wraps

def test_threaded():
    @threaded
    def function(expr, *args):
        return 2*expr + sum(args)

    assert function(Matrix([[x, y], [1, x]]), 1, 2) == \
        Matrix([[2*x+3, 2*y+3], [5, 2*x+3]])

    assert function(Eq(x, y), 1, 2) == Eq(2*x+3, 2*y+3)

    assert function([x, y], 1, 2) == [2*x+3, 2*y+3]
    assert function((x, y), 1, 2) == (2*x+3, 2*y+3)

    assert function(set([x, y]), 1, 2) == set([2*x+3, 2*y+3])

    @threaded
    def function(expr, n):
        return expr**n

    assert function(x + y, 2) == x**2 + y**2
    assert function(x, 2) == x**2

def test_xthreaded():
    @xthreaded
    def function(expr, n):
        return expr**n

    assert function(x + y, 2) == (x + y)**2

def test_wraps():
    def my_func(x):
        """My function. """

    my_func.is_my_func = True

    new_my_func = threaded(my_func)
    new_my_func = wraps(my_func)(new_my_func)

    assert new_my_func.__name__ == 'my_func'
    assert new_my_func.__doc__ == 'My function. '
    assert hasattr(new_my_func, 'is_my_func')
    assert new_my_func.is_my_func is True
