python_function_wrap.py

# coding: utf-8
"""Set the default arguments to functions.

>>> set_wrappers([module1, module2])  # doctest: +SKIP
{'f': OrderedDict([(u'module1.f.y', None), (u'module1.f.x', None), (u'a.f2.z', None)]}
>>> module1.fn(**{"module1.fn.arg1": 1, "module2.fn.arg1": 0}})  # doctest: +SKIP
module1.fn(**{"module1.fn.kwargs": {"args3": 1}}})  # doctest: +SKIP
"""
import importlib
import argparse
import inspect
import types
import functools
import collections


def get_default_args(fn):
    """
    :rtype: dict
    """
    spec = inspect.getargspec(fn)
    defaults = (spec.defaults if spec.defaults else [])
    fn_args = {i: None for i in spec.args[
        :len(spec.args) - len(defaults)]}
    default_kws = dict(zip(spec.args[-len(defaults):],
                           defaults))
    kws = collections.OrderedDict(fn_args.items() + default_kws.items())
    return kws


def wrapper(fn, module=None, *args, **kwargs):
    """
    :rtype: dict
    """
    path = u'.'.join(([module.__name__, fn.__name__]
                      if module else []))
    kwargs_path = path + u'.kwargs'
    spec = inspect.getargspec(fn)
    
    def _wrapper(*_args, **_kwargs):
        rest_args = spec.args[len(_args):]
        _kwargs = dict(_kwargs, **{
            k[len(path) + 1:]: v
            for k, v in _kwargs.items()
            if k.startswith(path)
            and k.split(u'.')[-1] in rest_args})
        if kwargs_path in _kwargs:
            _kwargs.update({
                k: v
                for k, v in _kwargs[kwargs_path].items()
                if k in rest_args})
        if spec.varargs and spec.keywords:
            return fn(*_args, **_kwargs)
        elif spec.varargs and not spec.keywords:
            return fn(*_args)
        elif not spec.varargs and spec.keywords:
            return fn(**_kwargs)
        elif not spec.varargs and not spec.keywords:
            return fn(*_args, **{
                k: v
                for k, v in _kwargs.items()
                if k in spec.args})
    return _wrapper


def get_deafault_function_args(module=None, module_path=None):
    """
    :rtype: dict
    """
    module_path, module = (
        (module_path, importlib.import_module(module_path))
        if module_path
        else (module.__name__, module))
    funcs = {
        k: v
        for k, v in vars(module).items()
        if isinstance(v, types.FunctionType)}
    infos = {}
    for name, func in funcs.items():
        spec = inspect.getargspec(func)
        defaults = (spec.defaults if spec.defaults else [])
        fn_args = [(i, None) for i in spec.args[
            :len(spec.args) - len(defaults)]]
        default_kws = zip(spec.args[-len(defaults):],
                          defaults)
        func_path = u'.'.join([module_path, name])
        kws = collections.OrderedDict(
            (fn_args + default_kws
             + ([(func_path + u'.args', [])] if spec.varargs else [])
             + ([(func_path + u'.kwargs', {})] if spec.keywords else [])))
        infos[name] = {
            u"paths": [
                (u'.'.join([func_path, k]), v)
                for k, v in kws.items()],
            u"args": kws.keys(),
            u"varargs": spec.varargs,
            u"keywords": spec.keywords,
        }
    return infos


def get_paths(obj):
    """
    :rtype: list

    >>> get_paths(get_deafault_function_args(os))  # doctest: +SKIP
    """
    # return list(itertools.chain.from_iterable([
    #     v['paths']
    #     for k, v in obj.items()]))
    return dict([(k, collections.OrderedDict(v[u'paths']))
                 for k, v in obj.items()])


def set_wrappers(modules):
    """
    :rtype: dict
    """
    values = reduce(lambda i, m: dict(i, **m),
                    [get_paths(get_deafault_function_args(module))
                     for module in modules])
    for module in modules:
        for k, v in vars(module).items():
            if isinstance(v, types.FunctionType):
                setattr(module, k, functools.wraps(v)(wrapper(v, module)))
    return values


def get_argparser(modules=[], parser=None):
    parser = (argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter)
              if not parser
              else parser)
    kwds = set_wrappers(modules)
    for k, v in kwds.items():
        for path, value in v.items():
            parser.add_argument(
                u'--' + path,
                default=value, help=' ')
    return parser


def test():
    import a, b
    print(set_wrappers([a, b]))
    print(a.f(2, **{"a.f.kwargs": {"x": 1}}))
    print(a.f3(2, **{"a.f.kwargs": {"x": 1}}))
    print(b.f4(2, **{"b.f4.kwargs": {"x": 1}}))