Skip to content

Commit

Permalink
pythongh-83076: ~4x speed improvement in (Async)Mock instantiation
Browse files Browse the repository at this point in the history
  • Loading branch information
carljm committed Dec 14, 2022
1 parent 9663853 commit 8a2a39f
Showing 1 changed file with 15 additions and 14 deletions.
29 changes: 15 additions & 14 deletions Lib/unittest/mock.py
Original file line number Diff line number Diff line change
Expand Up @@ -411,15 +411,18 @@ class NonCallableMock(Base):
# necessary.
_lock = RLock()

def __new__(cls, /, *args, **kw):
def __new__(
cls, spec=None, wraps=None, name=None, spec_set=None,
parent=None, _spec_state=None, _new_name='', _new_parent=None,
_spec_as_instance=False, _eat_self=None, unsafe=False, **kwargs
):
# every instance has its own class
# so we can create magic methods on the
# class without stomping on other mocks
bases = (cls,)
if not issubclass(cls, AsyncMockMixin):
# Check if spec is an async object or function
bound_args = _MOCK_SIG.bind_partial(cls, *args, **kw).arguments
spec_arg = bound_args.get('spec_set', bound_args.get('spec'))
spec_arg = spec_set or spec
if spec_arg is not None and _is_async_obj(spec_arg):
bases = (AsyncMockMixin, cls)
new = type(cls.__name__, bases, {'__doc__': cls.__doc__})
Expand Down Expand Up @@ -503,11 +506,6 @@ def _mock_add_spec(self, spec, spec_set, _spec_as_instance=False,

_spec_class = None
_spec_signature = None
_spec_asyncs = []

for attr in dir(spec):
if iscoroutinefunction(getattr(spec, attr, None)):
_spec_asyncs.append(attr)

if spec is not None and not _is_list(spec):
if isinstance(spec, type):
Expand All @@ -525,7 +523,6 @@ def _mock_add_spec(self, spec, spec_set, _spec_as_instance=False,
__dict__['_spec_set'] = spec_set
__dict__['_spec_signature'] = _spec_signature
__dict__['_mock_methods'] = spec
__dict__['_spec_asyncs'] = _spec_asyncs

def __get_return_value(self):
ret = self._mock_return_value
Expand Down Expand Up @@ -1015,7 +1012,8 @@ def _get_child_mock(self, /, **kw):
For non-callable mocks the callable variant will be used (rather than
any custom subclass)."""
_new_name = kw.get("_new_name")
if _new_name in self.__dict__['_spec_asyncs']:
_spec_val = getattr(self.__dict__["_spec_class"], _new_name, None)
if _spec_val is not None and asyncio.iscoroutinefunction(_spec_val):
return AsyncMock(**kw)

if self._mock_sealed:
Expand Down Expand Up @@ -1057,9 +1055,6 @@ def _calls_repr(self, prefix="Calls"):
return f"\n{prefix}: {safe_repr(self.mock_calls)}."


_MOCK_SIG = inspect.signature(NonCallableMock.__init__)


class _AnyComparer(list):
"""A list which checks if it contains a call which may have an
argument of ANY, flipping the components of item and self from
Expand Down Expand Up @@ -2183,6 +2178,10 @@ def __get__(self, obj, _type=None):
return self.create_mock()


_CODE_ATTRS = dir(CodeType)
_CODE_SIG = inspect.signature(partial(CodeType.__init__, None))


class AsyncMockMixin(Base):
await_count = _delegating_property('await_count')
await_args = _delegating_property('await_args')
Expand All @@ -2200,7 +2199,9 @@ def __init__(self, /, *args, **kwargs):
self.__dict__['_mock_await_count'] = 0
self.__dict__['_mock_await_args'] = None
self.__dict__['_mock_await_args_list'] = _CallList()
code_mock = NonCallableMock(spec_set=CodeType)
code_mock = NonCallableMock(spec_set=_CODE_ATTRS)
code_mock.__dict__["_spec_class"] = CodeType
code_mock.__dict__["_spec_signature"] = _CODE_SIG
code_mock.co_flags = inspect.CO_COROUTINE
self.__dict__['__code__'] = code_mock
self.__dict__['__name__'] = 'AsyncMock'
Expand Down

0 comments on commit 8a2a39f

Please sign in to comment.