Add kwargs from

easypy/decorations.py

@@ -2,7 +2,9 @@
 This module is about making it easier to create decorators
 """
 
+from collections import OrderedDict
 from functools import wraps, partial, update_wrapper
+from itertools import chain
 from operator import attrgetter
 from abc import ABCMeta, abstractmethod
 import inspect
@@ -188,3 +190,60 @@ def wrapper(func):
                 param_names=mismatches)
         return func
     return wrapper
+
+
+__KEYWORD_PARAMS = (inspect.Parameter.POSITIONAL_OR_KEYWORD, inspect.Parameter.KEYWORD_ONLY)
+
+
+def kwargs_from(*functions, exclude=()):
+    """
+    Edits the decorated function's signature to expand the variadic keyword
+    arguments parameter to the possible keywords from the wrapped functions.
+    This allows better completions inside interactive tools such as IPython.
+
+    :param functions: The functions to get the keywords from.
+    :param exclude: A list of parameters to exclude from the new signature.
+    :raises TypeError: When the decorated function does not have a variadic
+                       keyword argument.
+
+    >>> def foo(*, a, b, c):
+    ...     ...
+    >>> @kwargs_from(foo)
+    ... def bar(**kwargs):
+    ...     ...
+    >>> help(bar)
+    Help on function bar in module easypy.decorations:
+    <BLANKLINE>
+    bar(*, a, b, c)
+    <BLANKLINE>
+    """
+    exclude = set(exclude or ())
+    all_original_params = (inspect.signature(func).parameters for func in functions)
+    def _decorator(func):
+        signature = inspect.signature(func)
+
+        kws_param = None
+        params = OrderedDict()
+        for param in signature.parameters.values():
+            if param.kind != inspect.Parameter.VAR_KEYWORD:
+                params[param.name] = param
+            else:
+                kws_param = param
+        if kws_param is None:
+            raise TypeError("kwargs_from can only wrap functions with variadic keyword arguments")
+
+        keep_kwargs = False
+        for param in chain.from_iterable(original_params.values() for original_params in all_original_params):
+            if param.name in exclude:
+                pass
+            elif param.kind in __KEYWORD_PARAMS and param.name not in params:
+                params[param.name] = param.replace(kind=inspect.Parameter.KEYWORD_ONLY)
+            elif param.kind == inspect.Parameter.VAR_KEYWORD:
+                keep_kwargs = True
+
+        if keep_kwargs:
+            params['**'] = kws_param
+
+        func.__signature__ = signature.replace(parameters=params.values())
+        return func
+    return _decorator

tests/test_decorations.py

@@ -1,9 +1,12 @@
 import pytest
 
 from functools import wraps
+from io import StringIO
+from pydoc import doc
 
 from easypy.decorations import lazy_decorator
 from easypy.decorations import ensure_same_defaults, DefaultsMismatch
+from easypy.decorations import kwargs_from
 from easypy.misc import kwargs_resilient
 
 
@@ -190,3 +193,77 @@ def foo(a=1, b=3):
         def foo(a=2, b=3):
             pass
     assert exc.value.param_names == ['a']
+
+
+def signature_line(thing):
+    with StringIO() as capture:
+        doc(thing, output=capture)
+        return capture.getvalue().splitlines()[2]
+
+
+def test_kwargs_from():
+    import sys
+    if sys.version_info < (3, 7):
+        fix_annotations = lambda sign: sign.replace(': ', ':')
+    else:
+        fix_annotations = lambda sign: sign
+
+    def foo(*, a, b: int, c=3):
+        pass
+
+    @kwargs_from(foo)
+    def bar(a=1, **kwargs) -> bool:
+        return False
+
+    assert signature_line(bar) == fix_annotations('bar(a=1, *, b: int, c=3) -> bool')
+
+    @kwargs_from(bar)
+    def baz(c=4, **kwargs):
+        pass
+
+    assert signature_line(baz) == fix_annotations('baz(c=4, *, a=1, b: int)')
+
+
+def test_kwargs_from_no_kwargs():
+    with pytest.raises(TypeError):
+        @kwargs_from(lambda *, x, y: ...)
+        def foo(a):
+            pass
+
+    with pytest.raises(TypeError):
+        @kwargs_from(lambda *, x, y: ...)
+        def bar(x, y):
+            pass
+
+    @kwargs_from(lambda *args: ...)
+    def baz(a, **kwargs):
+        pass
+
+    assert signature_line(baz) == 'baz(a)'
+
+
+def test_kwargs_from_keep_kwargs():
+    def foo(*, a, b, **foos):
+        pass
+
+    @kwargs_from(foo)
+    def bar(x, **bars):
+        pass
+
+    assert signature_line(bar) == 'bar(x, *, a, b, **bars)'
+
+
+def test_kwargs_from_multiple():
+    @kwargs_from(lambda a, b: ..., lambda x, y, a: ...)
+    def bar(**kwargs):
+        pass
+
+    assert signature_line(bar) == 'bar(*, a, b, x, y)'
+
+
+def test_kwargs_from_exclude():
+    @kwargs_from(lambda a, b, c, w: ..., exclude=['w'])
+    def bar(**kwargs):
+        pass
+
+    assert signature_line(bar) == 'bar(*, a, b, c)'