"""
.. module:: factory
   :synopsis: Functions and decorators to construct nuts.
"""
from __future__ import absolute_import
import functools
from nutsflow.base import Nut, NutSink, NutSource, NutFunction
def _arg_insert(args, arg, pos=0):
    """
    Insert arg in args at given position.
    :param tuple args: Some function arguments
    :param any arg: Some function argument
    :param int pos: Insert position. If None argument is appended.
    :return: List with arguments where arg is inserted
    :rtype: list
    """
    args = list(args)
    if pos is None:
        args.append(arg)
    else:
        args.insert(pos, arg)
    return args
def _wrap(wrappercls, func):
    """
    Return wrapped function.
    Used to ensure that decorated nut function has the correct docstring.
    :param class wrappercls: Nut wrapper class
    :param function func: Function to wrap
    :return: Wrapped function
    :rtype: function
    """
    @functools.wraps(func)
    def wrapper(*args, **kwds):
        return wrappercls(*args, **kwds)
    return wrapper
def _create_nut_wrapper(base_class, func, iterpos):
    """
    Return Nut for given function.
    :param class base_class: Base class, e.g. Nut, NutSink, NutFunction, ...
    :param function func: Function to wrap
    :param int iterpos: Argument position for iterable in function.
    :return: Nut that wraps the given function.
    :rtype: Nut
    """
    class Wrapper(base_class):
        def __rrshift__(self, iterable):
            args = _arg_insert(self.args, iterable, iterpos)
            return func(*args, **self.kwargs)
    return _wrap(Wrapper, func)
def _create_filter_wrapper(func, invert=False):
    """
    Return filter Nut for given function.
    :param func: Filter function to wrap
    :param invert: Filter is inverted.
    :return: Nut operates as a filter.
    :rtype: Nut
    """
    class Wrapper(Nut):
        def __rrshift__(self, iterable):
            for e in iterable:
                args = _arg_insert(self.args, e)
                if bool(func(*args, **self.kwargs)) != invert:
                    yield e
    return _wrap(Wrapper, func)
[docs]def nut_function(func):
    """
    Decorator for Nut functions.
    Example on how to define a custom function nut:
    .. code::
      @nut_function
      def TimesN(x, n):
          return x * n
      [1, 2, 3] >> TimesN(2) >> Collect()  -->  [2, 4, 6]
    :param function func: Function to decorate
    :return: Nut function for given function
    :rtype: NutFunction
    """
    class Wrapper(NutFunction):
        def __call__(self, element):
            return func(element, *self.args, **self.kwargs)
    return _wrap(Wrapper, func) 
[docs]def nut_source(func):
    """
    Decorator for Nut sources.
    Example on how to define a custom source nut. Note that a source
    must return an iterable/generator and does not read any input.
    .. code::
      @nut_source
      def MyRange(start, end):
          return range(start, end)
      MyRange(0, 5) >> Collect()  --> [0, 1, 2, 3, 4]
    .. code::
      @nut_source
      def MyRange2(start, end):
          for i in range(start, end):
              yield i * 2
      MyRange2(0, 5) >> Collect()  --> [0, 2, 4, 6, 8]
    :param function func: Function to decorate
    :return: Nut source for given function
    :rtype: NutSource
    """
    class Wrapper(NutSource):
        def __iter__(self):
            return func(*self.args, **self.kwargs)
    return _wrap(Wrapper, func) 
[docs]def nut_processor(func, iterpos=0):
    """
    Decorator for Nut processors.
    Examples on how to define a custom processor nut.
    Note that a processor reads an iterable and must return
    an iterable/generator
    .. code::
      @nut_processor
      def Twice(iterable):
          for e in iterable:
              yield e
              yield e
      [1, 2, 3] >> Twice() >> Collect()  --> [1, 1, 2, 2, 3, 3]
    .. code::
      @nut_processor
      def Odd(iterable):
          return (e for e in iterable if e % 2)
      [1, 2, 3, 4, 5] >> Odd() >> Collect()  --> [1, 3, 5]
    .. code::
      @nut_processor
      def Clone(iterable, n):
          for e in iterable:
              for _ in range(p):
                  yield e
      [1, 2, 3] >> Clone(2) >> Collect()  --> [1, 1, 2, 2, 3, 3]
    :param function func: Function to decorate
    :param iterpos: Position of iterable in function arguments
    :return: Nut processor for given function
    :rtype: Nut
    """
    return _create_nut_wrapper(Nut, func, iterpos) 
[docs]def nut_sink(func, iterpos=0):
    """
    Decorator for Nut sinks.
    Example on how to define a custom sink nut:
    .. code::
      @nut_sink
      def ToList(iterable):
          return list(iterable)
      range(5) >> ToList()  -->   [0, 1, 2, 3, 4]
    .. code::
      @nut_sink
      def MyCollect(iterable, container):
          return container(iterable)
      range(5) >> MyCollect(tuple)  -->   (0, 1, 2, 3, 4)
    .. code::
      @nut_sink
      def MyProd(iterable):
          p = 1
          for e in iterable:
              p *= e
          return p
      [1, 2, 3] >> MyProd()  --> 12
    :param function func: Function to decorate
    :param iterpos: Position of iterable in function arguments
    :return: Nut sink for given function
    :rtype: NutSink
    """
    return _create_nut_wrapper(NutSink, func, iterpos) 
[docs]def nut_filter(func):
    """
    Decorator for Nut filters.
    Also see nut_filerfalse().
    Example on how to define a custom filter nut:
    .. code::
      @nut_filter
      def Positive(x):
          return x > 0
      [-1, 1, -2, 2] >> Positive() >> Collect()  --> [1, 2]
    .. code::
      @nut_filter
      def GreaterThan(x, threshold):
          return x > threshold
      [1, 2, 3, 4] >> GreaterThan(2) >> Collect()  --> [3, 4] 
    :param function func: Function to decorate. Must return boolean value.
    :return: Nut filter for given function
    :rtype: Nut
    """
    return _create_filter_wrapper(func, invert=False) 
[docs]def nut_filterfalse(func):
    """
    Decorator for Nut filters that are inverted.
    Also see nut_filter().
    Example on how to define a custom filter-false nut:
    .. code::
      @nut_filterfalse
      def NotGreaterThan(x, threshold):
          return x > threshold
      [1, 2, 3, 4] >> NotGreaterThan(2) >> Collect()  --> [1, 2]
    :param function func: Function to decorate
    :return: Nut filter for given function. . Must return boolean value.
    :rtype: Nut
    """
    return _create_filter_wrapper(func, invert=True)