Use annotations in autospec?

Ideas
1

The unittest.mock autospeccing feature is handy; mocks are constrained to only allow attributes and method calls that exist on a target object. There seem to be two known limitations that I think could be addressed by leveraging type annotations:

  • Types of return values are not known, so as soon as you make any chained method calls the autospeccing goes away
  • As the docs call out:

    A more serious problem is that it is common for instance attributes to be
    created in the __init__()method and not to exist on the class at all.
    autospec can’t know about any dynamically created attributes and restricts
    the api to visible attributes.

Both of these can be addressed by annotations:

  • if a callable has a return type annotation, autospeccing can use that
  • if an instance variable annotation exists for an an attribute created in __init__(), autospeccing can use that

I took a few minutes to check out the feasibility of this and it seems pretty doable at least for basic cases where the annotation is a concrete type. But there are probably quite a few edge cases to consider.

Is anyone else interested in having or implementing this?

diff --git a/Lib/test/test_unittest/testmock/testhelpers.py b/Lib/test/test_unittest/testmock/testhelpers.py
index 9e7ec5d62d5d..e4743631f17d 100644
--- a/Lib/test/test_unittest/testmock/testhelpers.py
+++ b/Lib/test/test_unittest/testmock/testhelpers.py
@@ -618,6 +618,12 @@ class Sub(SomeClass):
            self.assertRaises(AttributeError, setattr, mock, 'foo', 'bar')
            self.assertRaises(AttributeError, setattr, mock.attr, 'foo', 'bar')

+    def test_return_value_from_annotation(self):
+        def fn() -> str:
+            return "Hello world"
+
+        mock = create_autospec(fn)
+        self.assertIsInstance(mock(), str)

    def test_descriptors(self):
        class Foo(object):
diff --git a/Lib/unittest/mock.py b/Lib/unittest/mock.py
index 0f93cb53c3d5..cabb55425d48 100644
--- a/Lib/unittest/mock.py
+++ b/Lib/unittest/mock.py
@@ -2736,6 +2736,11 @@ def create_autospec(spec, spec_set=False, instance=False, _parent=None,
        mock = _set_signature(mock, spec)
        if is_async_func:
            _setup_async_mock(mock)
+        if 'return_value' not in kwargs:
+            annotations = inspect.get_annotations(spec)
+            if 'return' in annotations:
+                mock.return_value = create_autospec(annotations['return'])
+                # What about list[str]? union types?
    else:
        _check_signature(spec, mock, is_type, instance)
1 Like