bpo-36996: Handle async functions when mock.patch is used as a decorator

[tirkarthi]

Return a coroutine while patching async functions with a decorator.

Co-authored-by: Andrew Svetlov [email protected]

https://bugs.python.org/issue36996

[tirkarthi]

I have tried using below pattern but I am not sure modifying the return value data type based on a flag as a good idea. Also using yield from makes it not usable as a function with the return statement in else part since patched is marked as a generator. So I have used a separate decorate_async_callable as suggested and my concern is that future changes to decorate_callable need to be done in decorate_async_callable too due to the code duplication. If there is any better way to handle this it would be great to know.

def decorate_callable(self, func, async_func=False):
    def patched(func):
        if async_func:
            yield from func(*args, **keywargs)
        else:
            return func(*args, **keywargs)
    if async_func:
        patched = types.coroutine(async_func)  

cc: @lisroach @asvetlov

[asvetlov]

Interesting challenge.
Let me try to find a way to don't duplicate too much.

[asvetlov]

Please replace docstring with just a comment.
Docstrings are for public API but this class is entirely an implementation detail.

In the comment please add an explicit text like NB. Keep the method in sync with decorate_callable.
Add the same comment to decorate_callable.

[asvetlov]

This seems to work fine:

    @contextlib.contextmanager
    def decoration_helper(self, patched, args, keywargs):
        extra_args = []
        entered_patchers = []
        patching = None

        exc_info = tuple()
        try:
            for patching in patched.patchings:
                arg = patching.__enter__()
                entered_patchers.append(patching)
                if patching.attribute_name is not None:
                    keywargs.update(arg)
                elif patching.new is DEFAULT:
                    extra_args.append(arg)

            args += tuple(extra_args)
            yield (args, keywargs)
        except:
            if (patching not in entered_patchers and
                _is_started(patching)):
                # the patcher may have been started, but an exception
                # raised whilst entering one of its additional_patchers
                entered_patchers.append(patching)
            # Pass the exception to __exit__
            exc_info = sys.exc_info()
            # re-raise the exception
            raise
        finally:
            for patching in reversed(entered_patchers):
                patching.__exit__(*exc_info)


    def decorate_callable(self, func):
        if hasattr(func, 'patchings'):
            func.patchings.append(self)
            return func

        @wraps(func)
        def patched(*args, **keywargs):
            with self.decoration_helper(patched,
                                        args,
                                        keywargs) as (newargs, newkeywargs):
                return func(*newargs, **newkeywargs)

        patched.patchings = [self]
        return patched


    def decorate_async_callable(self, func):
        '''This is same as decorate_callable except patched is a coroutine'''
        if hasattr(func, 'patchings'):
            func.patchings.append(self)
            return func

        @wraps(func)
        async def patched(*args, **keywargs):
            with self.decoration_helper(patched,
                                        args,
                                        keywargs) as (newargs, newkeywargs):
                return await func(*newargs, **newkeywargs)

        patched.patchings = [self]
        return patched

[tirkarthi]

This looks like a cool pattern to me on the control flow with yield and acts like the code is pasted except the logic can be used for sync and async functions.

[asvetlov]

Looks good, a few very minor comments

[asvetlov]

thanks!

[tirkarthi]

Thanks @asvetlov for the review and merge.