Skip to content

Commit c5726b7

Browse files
authoredDec 23, 2022
gh-83076: 3.8x speed improvement in (Async)Mock instantiation (#100252)
1 parent a98d9ea commit c5726b7

File tree

3 files changed

+36
-16
lines changed

3 files changed

+36
-16
lines changed
 

‎Lib/test/test_unittest/testmock/testasync.py

+13
Original file line numberDiff line numberDiff line change
@@ -300,6 +300,19 @@ def test_spec_normal_methods_on_class_with_mock(self):
300300
self.assertIsInstance(mock.async_method, AsyncMock)
301301
self.assertIsInstance(mock.normal_method, Mock)
302302

303+
def test_spec_async_attributes_instance(self):
304+
async_instance = AsyncClass()
305+
async_instance.async_func_attr = async_func
306+
async_instance.later_async_func_attr = normal_func
307+
308+
mock_async_instance = Mock(spec_set=async_instance)
309+
310+
async_instance.later_async_func_attr = async_func
311+
312+
self.assertIsInstance(mock_async_instance.async_func_attr, AsyncMock)
313+
# only the shape of the spec at the time of mock construction matters
314+
self.assertNotIsInstance(mock_async_instance.later_async_func_attr, AsyncMock)
315+
303316
def test_spec_mock_type_kw(self):
304317
def inner_test(mock_type):
305318
async_mock = mock_type(spec=async_func)

‎Lib/unittest/mock.py

+22-16
Original file line numberDiff line numberDiff line change
@@ -411,15 +411,18 @@ class NonCallableMock(Base):
411411
# necessary.
412412
_lock = RLock()
413413

414-
def __new__(cls, /, *args, **kw):
414+
def __new__(
415+
cls, spec=None, wraps=None, name=None, spec_set=None,
416+
parent=None, _spec_state=None, _new_name='', _new_parent=None,
417+
_spec_as_instance=False, _eat_self=None, unsafe=False, **kwargs
418+
):
415419
# every instance has its own class
416420
# so we can create magic methods on the
417421
# class without stomping on other mocks
418422
bases = (cls,)
419423
if not issubclass(cls, AsyncMockMixin):
420424
# Check if spec is an async object or function
421-
bound_args = _MOCK_SIG.bind_partial(cls, *args, **kw).arguments
422-
spec_arg = bound_args.get('spec_set', bound_args.get('spec'))
425+
spec_arg = spec_set or spec
423426
if spec_arg is not None and _is_async_obj(spec_arg):
424427
bases = (AsyncMockMixin, cls)
425428
new = type(cls.__name__, bases, {'__doc__': cls.__doc__})
@@ -505,10 +508,6 @@ def _mock_add_spec(self, spec, spec_set, _spec_as_instance=False,
505508
_spec_signature = None
506509
_spec_asyncs = []
507510

508-
for attr in dir(spec):
509-
if iscoroutinefunction(getattr(spec, attr, None)):
510-
_spec_asyncs.append(attr)
511-
512511
if spec is not None and not _is_list(spec):
513512
if isinstance(spec, type):
514513
_spec_class = spec
@@ -518,7 +517,13 @@ def _mock_add_spec(self, spec, spec_set, _spec_as_instance=False,
518517
_spec_as_instance, _eat_self)
519518
_spec_signature = res and res[1]
520519

521-
spec = dir(spec)
520+
spec_list = dir(spec)
521+
522+
for attr in spec_list:
523+
if iscoroutinefunction(getattr(spec, attr, None)):
524+
_spec_asyncs.append(attr)
525+
526+
spec = spec_list
522527

523528
__dict__ = self.__dict__
524529
__dict__['_spec_class'] = _spec_class
@@ -1057,9 +1062,6 @@ def _calls_repr(self, prefix="Calls"):
10571062
return f"\n{prefix}: {safe_repr(self.mock_calls)}."
10581063

10591064

1060-
_MOCK_SIG = inspect.signature(NonCallableMock.__init__)
1061-
1062-
10631065
class _AnyComparer(list):
10641066
"""A list which checks if it contains a call which may have an
10651067
argument of ANY, flipping the components of item and self from
@@ -2138,10 +2140,8 @@ def mock_add_spec(self, spec, spec_set=False):
21382140

21392141

21402142
class AsyncMagicMixin(MagicMixin):
2141-
def __init__(self, /, *args, **kw):
2142-
self._mock_set_magics() # make magic work for kwargs in init
2143-
_safe_super(AsyncMagicMixin, self).__init__(*args, **kw)
2144-
self._mock_set_magics() # fix magic broken by upper level init
2143+
pass
2144+
21452145

21462146
class MagicMock(MagicMixin, Mock):
21472147
"""
@@ -2183,6 +2183,10 @@ def __get__(self, obj, _type=None):
21832183
return self.create_mock()
21842184

21852185

2186+
_CODE_ATTRS = dir(CodeType)
2187+
_CODE_SIG = inspect.signature(partial(CodeType.__init__, None))
2188+
2189+
21862190
class AsyncMockMixin(Base):
21872191
await_count = _delegating_property('await_count')
21882192
await_args = _delegating_property('await_args')
@@ -2200,7 +2204,9 @@ def __init__(self, /, *args, **kwargs):
22002204
self.__dict__['_mock_await_count'] = 0
22012205
self.__dict__['_mock_await_args'] = None
22022206
self.__dict__['_mock_await_args_list'] = _CallList()
2203-
code_mock = NonCallableMock(spec_set=CodeType)
2207+
code_mock = NonCallableMock(spec_set=_CODE_ATTRS)
2208+
code_mock.__dict__["_spec_class"] = CodeType
2209+
code_mock.__dict__["_spec_signature"] = _CODE_SIG
22042210
code_mock.co_flags = inspect.CO_COROUTINE
22052211
self.__dict__['__code__'] = code_mock
22062212
self.__dict__['__name__'] = 'AsyncMock'
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
Instantiation of ``Mock()`` and ``AsyncMock()`` is now 3.8x faster.

0 commit comments

Comments
 (0)