gh-98086: Now `patch.dict` can decorate async functions (#98095)

This commit is contained in:
Nikita Sobolev 2022-11-11 11:04:30 +03:00 committed by GitHub
parent 97c493dd35
commit 67b4d2772c
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
3 changed files with 36 additions and 0 deletions

View file

@ -1809,6 +1809,12 @@ class _patch_dict(object):
def __call__(self, f):
if isinstance(f, type):
return self.decorate_class(f)
if inspect.iscoroutinefunction(f):
return self.decorate_async_callable(f)
return self.decorate_callable(f)
def decorate_callable(self, f):
@wraps(f)
def _inner(*args, **kw):
self._patch_dict()
@ -1820,6 +1826,18 @@ class _patch_dict(object):
return _inner
def decorate_async_callable(self, f):
@wraps(f)
async def _inner(*args, **kw):
self._patch_dict()
try:
return await f(*args, **kw)
finally:
self._unpatch_dict()
return _inner
def decorate_class(self, klass):
for attr in dir(klass):
attr_value = getattr(klass, attr)