Example #1
0
def update_wrapper(
    wrapper: T.Callable,
    wrapped: T.Callable,
    signature: T.Union[_FullerSig, None, bool] = True,  # not in functools
    docstring: T.Union[str, bool] = True,  # not in functools
    assigned: T.Sequence[str] = WRAPPER_ASSIGNMENTS,
    updated: T.Sequence[str] = WRAPPER_UPDATES,
    # docstring options
    _doc_fmt: T.Optional[dict] = None,  # not in functools
    _doc_style: T.Union[str, T.Callable, None] = None,
):
    """Update a wrapper function to look like the wrapped function.

    Parameters
    ----------
    wrapper : Callable
        the function to be updated
    wrapped : Callable
       the original function
    signature : Signature or None or bool, optional
        signature to impose on `wrapper`.
        None and False default to `wrapped`'s signature.
        True merges `wrapper` and `wrapped` kwdefaults & annotations
    docstring : str or bool, optional
        docstring to impose on `wrapper`.
        False ignores `wrapper`'s docstring, using only `wrapped`'s docstring.
        None (defualt) merges the `wrapper` and `wrapped` docstring
    assigned : tuple, optional
       tuple naming the attributes assigned directly
       from the wrapped function to the wrapper function (defaults to
       ``functools.WRAPPER_ASSIGNMENTS``)
    updated : tuple, optional
       is a tuple naming the attributes of the wrapper that
       are updated with the corresponding attribute from the wrapped
       function (defaults to ``functools.WRAPPER_UPDATES``)
    _doc_fmt : dict, optional
        dictionary to format wrapper docstring
    _doc_style: str or Callable, optional
        the style of the docstring
        if None (default), appends `wrapper` docstring
        if str or Callable, merges the docstring

    Returns
    -------
    wrapper : Callable
        `wrapper` function updated by the `wrapped` function's attributes and
        also the provided `signature` and `docstring`.

    Raises
    ------
    ValueError
        if docstring is True

    """
    # ---------------------------------------
    # preamble

    signature, _update_sig = __parse_sig_for_update_wrapper(signature, wrapped)

    # need to get wrapper properties now
    wrapper_sig = _FullerSig.from_callable(wrapper)

    wrapper_doc = _nspct.getdoc(wrapper) or ""
    wrapper_doc = "\n".join(wrapper_doc.split("\n")[1:])  # drop title

    if _doc_fmt is None:
        _doc_fmt = {}

    # ---------------------------------------
    # update wrapper (same as functools.update_wrapper)

    for attr in assigned:
        try:
            value = getattr(wrapped, attr)
        except AttributeError:
            pass
        else:
            setattr(wrapper, attr, value)

    for attr in updated:  # update whole dictionary
        getattr(wrapper, attr).update(getattr(wrapped, attr, {}))

    # ---------------------------------------

    # deal with signature
    if signature in (None, False):
        pass

    elif _update_sig:  # merge wrapped and wrapper signature

        signature = __update_wrapper_update_sig(
            signature, wrapper_sig, _doc_fmt
        )

        for attr in SIGNATURE_ASSIGNMENTS:
            value = getattr(signature, attr)
            setattr(wrapper, attr, value)

        wrapper.__signature__ = signature.signature

    else:  # a signature object
        for attr in SIGNATURE_ASSIGNMENTS:
            _value = getattr(signature, attr)
            setattr(wrapper, attr, _value)

        # for docstring
        for param in wrapper_sig.parameters.values():
            # can only merge keyword-only
            if param.kind == _nspct.KEYWORD_ONLY:
                _doc_fmt[param.name] = param.default

        wrapper.__signature__ = signature.signature

    # ---------------------------------------
    # docstring

    if _doc_fmt:  # (not empty dict)
        wrapper_doc = _FormatTemplate(wrapper_doc).safe_substitute(**_doc_fmt)

    wrapper.__doc__ = __update_wrapper_docstring(
        wrapped,
        docstring=docstring,
        wrapper_doc=wrapper_doc,
        _doc_style=_doc_style,
    )

    # Issue #17482: set __wrapped__ last so we don't inadvertently copy it
    # from the wrapped function when updating __dict__
    wrapper.__wrapped__ = wrapped
    # Return the wrapper so this can be used as a decorator via partial()
    return wrapper