gh-124176: Add special support for dataclasses to create_autospec (#124429)

Backports: 3a0e7f57628466aedcaaf6c5ff7c8224f5155a2c
Signed-off-by: Chris Withers <[email protected]>

Author: sobolevn

NEWS.d/2024-09-24-13-32-16.gh-issue-124176.6hmOPz.rst

@@ -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.

mock/mock.py

@@ -32,6 +32,13 @@
 import sys
 import threading
 import builtins
+
+try:
+    from dataclasses import fields, is_dataclass
+    HAS_DATACLASSES = True
+except ImportError:
+    HAS_DATACLASSES = False
+
 from types import CodeType, ModuleType, MethodType
 from unittest.util import safe_repr
 from functools import wraps, partial
@@ -2808,7 +2815,15 @@ 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}
+
+    entries = [(entry, _missing) for entry in dir(spec)]
+    if is_type and instance and HAS_DATACLASSES and is_dataclass(spec):
+        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:
@@ -2865,7 +2880,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
@@ -2879,10 +2894,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 _missing:
+            try:
+                original = getattr(spec, entry)
+            except AttributeError:
+                continue
 
         child_kwargs = {'spec': original}
         # Wrap child attributes also.

mock/tests/testhelpers.py

@@ -10,10 +10,17 @@
 from mock.mock import _Call, _CallList, _callable
 from mock import IS_PYPY
 
+try:
+    from dataclasses import dataclass, field, InitVar
+except ImportError:
+    pass
+
 from datetime import datetime
 from functools import partial
+from typing import ClassVar
 
 import pytest
+import sys
 
 
 class SomeClass(object):
@@ -1043,6 +1050,80 @@ def f(a): pass
         self.assertEqual(mock.mock_calls, [])
         self.assertEqual(rv.mock_calls, [])
 
+    @pytest.mark.skipif(sys.version_info < (3, 7), reason="requires python 3.7 or higher")
+    def test_dataclass_post_init(self):
+        @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)
+
+        # Classes do not have these fields:
+        mock = create_autospec(WithPostInit)
+        msg = "Mock object has no attribute"
+        with self.assertRaisesRegex(AttributeError, msg):
+            mock.a
+        with self.assertRaisesRegex(AttributeError, msg):
+            mock.b
+
+    @pytest.mark.skipif(sys.version_info < (3, 7), reason="requires python 3.7 or higher")
+    def test_dataclass_default(self):
+        @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)
+
+    @pytest.mark.skipif(sys.version_info < (3, 7), reason="requires python 3.7 or higher")
+    def test_dataclass_with_method(self):
+        @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()
+
+    @pytest.mark.skipif(sys.version_info < (3, 7), reason="requires python 3.7 or higher")
+    def test_dataclass_with_non_fields(self):
+        @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):