Skip to content

Commit

Permalink
pythongh-124176: Add special support for dataclasses to `create_autos…
Browse files Browse the repository at this point in the history
…pec`
  • Loading branch information
sobolevn committed Sep 24, 2024
1 parent e670a11 commit 312e880
Show file tree
Hide file tree
Showing 3 changed files with 85 additions and 6 deletions.
62 changes: 62 additions & 0 deletions Lib/test/test_unittest/testmock/testhelpers.py
Original file line number Diff line number Diff line change
Expand Up @@ -1034,6 +1034,68 @@ def f(a): pass
self.assertEqual(mock.mock_calls, [])
self.assertEqual(rv.mock_calls, [])

def test_dataclass(self):
from dataclasses import dataclass, field, InitVar
from typing import ClassVar

@dataclass
class WithPostInit:
a: int = field(init=False)
b: int = field(init=False)
def __post_init__(self):
self.a = 1
self.b = 2

for mock in [
create_autospec(WithPostInit, instance=True),
create_autospec(WithPostInit()),
]:
with self.subTest(mock=mock):
self.assertIsInstance(mock.a, int)
self.assertIsInstance(mock.b, int)

@dataclass
class WithDefault:
a: int
b: int = 0

for mock in [
create_autospec(WithDefault, instance=True),
create_autospec(WithDefault(1)),
]:
with self.subTest(mock=mock):
self.assertIsInstance(mock.a, int)
self.assertIsInstance(mock.b, int)

@dataclass
class WithMethod:
a: int
def b(self) -> int:
return 1

for mock in [
create_autospec(WithMethod, instance=True),
create_autospec(WithMethod(1)),
]:
with self.subTest(mock=mock):
self.assertIsInstance(mock.a, int)
mock.b.assert_not_called()

@dataclass
class WithNonFields:
a: ClassVar[int]
b: InitVar[int]

msg = "Mock object has no attribute"
for mock in [
create_autospec(WithNonFields, instance=True),
create_autospec(WithNonFields(1)),
]:
with self.subTest(mock=mock):
with self.assertRaisesRegex(AttributeError, msg):
mock.a
with self.assertRaisesRegex(AttributeError, msg):
mock.b

class TestCallList(unittest.TestCase):

Expand Down
25 changes: 19 additions & 6 deletions Lib/unittest/mock.py
Original file line number Diff line number Diff line change
Expand Up @@ -2754,7 +2754,19 @@ def create_autospec(spec, spec_set=False, instance=False, _parent=None,
raise InvalidSpecError(f'Cannot autospec a Mock object. '
f'[object={spec!r}]')
is_async_func = _is_async_func(spec)
_kwargs = {'spec': spec}

placeholder = object()
entries = [(entry, placeholder) for entry in dir(spec)]
# Not using `is_dataclass` to avoid an import of dataclasses module
# for types that don't need that.
if is_type and hasattr(spec, '__dataclass_fields__'):
from dataclasses import fields
dataclass_fields = fields(spec)
entries.extend((f.name, f.type) for f in dataclass_fields)
_kwargs = {'spec': [f.name for f in dataclass_fields]}
else:
_kwargs = {'spec': spec}

if spec_set:
_kwargs = {'spec_set': spec}
elif spec is None:
Expand Down Expand Up @@ -2811,7 +2823,7 @@ def create_autospec(spec, spec_set=False, instance=False, _parent=None,
_name='()', _parent=mock,
wraps=wrapped)

for entry in dir(spec):
for entry, original in entries:
if _is_magic(entry):
# MagicMock already does the useful magic methods for us
continue
Expand All @@ -2825,10 +2837,11 @@ def create_autospec(spec, spec_set=False, instance=False, _parent=None,
# AttributeError on being fetched?
# we could be resilient against it, or catch and propagate the
# exception when the attribute is fetched from the mock
try:
original = getattr(spec, entry)
except AttributeError:
continue
if original is placeholder:
try:
original = getattr(spec, entry)
except AttributeError:
continue

child_kwargs = {'spec': original}
# Wrap child attributes also.
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,4 @@
Add support for :func:`dataclasses.dataclass` in
:func:`unittest.mock.create_autospec`. Now ``create_autospec`` will check
for potential dataclasses and use :func:`dataclasses.fields` function to
retrieve the spec information.

0 comments on commit 312e880

Please sign in to comment.