Close #13266: Add inspect.unwrap

Initial patch by Daniel Urban and Aaron Iles

5 files changed, 141 insertions, 11 deletions

diff --git a/Doc/library/inspect.rst b/Doc/library/inspect.rst
index 40f482b..af6c96b 100644
--- a/Doc/library/inspect.rst
+++ b/Doc/library/inspect.rst
@@ -797,6 +797,23 @@ Classes and functions
    .. versionadded:: 3.3
 
 
+.. function:: unwrap(func, *, stop=None)
+
+   Get the object wrapped by *func*. It follows the chain of :attr:`__wrapped__`
+   attributes returning the last object in the chain.
+
+   *stop* is an optional callback accepting an object in the wrapper chain
+   as its sole argument that allows the unwrapping to be terminated early if
+   the callback returns a true value. If the callback never returns a true
+   value, the last object in the chain is returned as usual. For example,
+   :func:`signature` uses this to stop unwrapping if any object in the
+   chain has a ``__signature__`` attribute defined.
+
+   :exc:`ValueError` is raised if a cycle is encountered.
+
+   .. versionadded:: 3.4
+
+
 .. _inspect-stack:
 
 The interpreter stack
diff --git a/Doc/whatsnew/3.4.rst b/Doc/whatsnew/3.4.rst
index 40b8243..b5be568 100644
--- a/Doc/whatsnew/3.4.rst
+++ b/Doc/whatsnew/3.4.rst
@@ -185,6 +185,15 @@ functools
 
 New :func:`functools.singledispatch` decorator: see the :pep:`443`.
 
+
+inspect
+-------
+
+:func:`~inspect.unwrap` makes it easy to unravel wrapper function chains
+created by :func:`functools.wraps` (and any other API that sets the
+``__wrapped__`` attribute on a wrapper function).
+
+
 smtplib
 -------
 
@@ -327,6 +336,5 @@ that may require changes to your code.
   wrapped attribute set. This means ``__wrapped__`` attributes now correctly
   link a stack of decorated functions rather than every ``__wrapped__``
   attribute in the chain referring to the innermost function. Introspection
-  libraries that assumed the previous behaviour was intentional will need to
-  be updated to walk the chain of ``__wrapped__`` attributes to find the
-  innermost function.
+  libraries that assumed the previous behaviour was intentional can use
+  :func:`inspect.unwrap` to gain equivalent behaviour.
diff --git a/Lib/inspect.py b/Lib/inspect.py
index 4a28507..195c9fd 100644
--- a/Lib/inspect.py
+++ b/Lib/inspect.py
@@ -360,6 +360,40 @@ def getmro(cls):
     "Return tuple of base classes (including cls) in method resolution order."
     return cls.__mro__
 
+# -------------------------------------------------------- function helpers
+
+def unwrap(func, *, stop=None):
+    """Get the object wrapped by *func*.
+
+   Follows the chain of :attr:`__wrapped__` attributes returning the last
+   object in the chain.
+
+   *stop* is an optional callback accepting an object in the wrapper chain
+   as its sole argument that allows the unwrapping to be terminated early if
+   the callback returns a true value. If the callback never returns a true
+   value, the last object in the chain is returned as usual. For example,
+   :func:`signature` uses this to stop unwrapping if any object in the
+   chain has a ``__signature__`` attribute defined.
+
+   :exc:`ValueError` is raised if a cycle is encountered.
+
+    """
+    if stop is None:
+        def _is_wrapper(f):
+            return hasattr(f, '__wrapped__')
+    else:
+        def _is_wrapper(f):
+            return hasattr(f, '__wrapped__') and not stop(f)
+    f = func  # remember the original func for error reporting
+    memo = {id(f)} # Memoise by id to tolerate non-hashable objects
+    while _is_wrapper(func):
+        func = func.__wrapped__
+        id_func = id(func)
+        if id_func in memo:
+            raise ValueError('wrapper loop when unwrapping {!r}'.format(f))
+        memo.add(id_func)
+    return func
+
 # -------------------------------------------------- source code extraction
 def indentsize(line):
     """Return the indent size, in spaces, at the start of a line of text."""
@@ -1346,6 +1380,9 @@ def signature(obj):
         sig = signature(obj.__func__)
         return sig.replace(parameters=tuple(sig.parameters.values())[1:])
 
+    # Was this function wrapped by a decorator?
+    obj = unwrap(obj, stop=(lambda f: hasattr(f, "__signature__")))
+
     try:
         sig = obj.__signature__
     except AttributeError:
@@ -1354,13 +1391,6 @@ def signature(obj):
         if sig is not None:
             return sig
 
-    try:
-        # Was this function wrapped by a decorator?
-        wrapped = obj.__wrapped__
-    except AttributeError:
-        pass
-    else:
-        return signature(wrapped)
 
     if isinstance(obj, types.FunctionType):
         return Signature.from_function(obj)
diff --git a/Lib/test/test_inspect.py b/Lib/test/test_inspect.py
index 6bd9bd1..5de6212 100644
--- a/Lib/test/test_inspect.py
+++ b/Lib/test/test_inspect.py
@@ -8,6 +8,7 @@ import datetime
 import collections
 import os
 import shutil
+import functools
 from os.path import normcase
 
 from test.support import run_unittest, TESTFN, DirsOnSysPath
@@ -1719,6 +1720,17 @@ class TestSignatureObject(unittest.TestCase):
                          ((('b', ..., ..., "positional_or_keyword"),),
                           ...))
 
+        # Test we handle __signature__ partway down the wrapper stack
+        def wrapped_foo_call():
+            pass
+        wrapped_foo_call.__wrapped__ = Foo.__call__
+
+        self.assertEqual(self.signature(wrapped_foo_call),
+                         ((('a', ..., ..., "positional_or_keyword"),
+                           ('b', ..., ..., "positional_or_keyword")),
+                          ...))
+
+
     def test_signature_on_class(self):
         class C:
             def __init__(self, a):
@@ -1833,6 +1845,10 @@ class TestSignatureObject(unittest.TestCase):
         self.assertEqual(self.signature(Wrapped),
                          ((('a', ..., ..., "positional_or_keyword"),),
                           ...))
+        # wrapper loop:
+        Wrapped.__wrapped__ = Wrapped
+        with self.assertRaisesRegex(ValueError, 'wrapper loop'):
+            self.signature(Wrapped)
 
     def test_signature_on_lambdas(self):
         self.assertEqual(self.signature((lambda a=10: a)),
@@ -2284,6 +2300,62 @@ class TestBoundArguments(unittest.TestCase):
         self.assertNotEqual(ba, ba4)
 
 
+class TestUnwrap(unittest.TestCase):
+
+    def test_unwrap_one(self):
+        def func(a, b):
+            return a + b
+        wrapper = functools.lru_cache(maxsize=20)(func)
+        self.assertIs(inspect.unwrap(wrapper), func)
+
+    def test_unwrap_several(self):
+        def func(a, b):
+            return a + b
+        wrapper = func
+        for __ in range(10):
+            @functools.wraps(wrapper)
+            def wrapper():
+                pass
+        self.assertIsNot(wrapper.__wrapped__, func)
+        self.assertIs(inspect.unwrap(wrapper), func)
+
+    def test_stop(self):
+        def func1(a, b):
+            return a + b
+        @functools.wraps(func1)
+        def func2():
+            pass
+        @functools.wraps(func2)
+        def wrapper():
+            pass
+        func2.stop_here = 1
+        unwrapped = inspect.unwrap(wrapper,
+                                   stop=(lambda f: hasattr(f, "stop_here")))
+        self.assertIs(unwrapped, func2)
+
+    def test_cycle(self):
+        def func1(): pass
+        func1.__wrapped__ = func1
+        with self.assertRaisesRegex(ValueError, 'wrapper loop'):
+            inspect.unwrap(func1)
+
+        def func2(): pass
+        func2.__wrapped__ = func1
+        func1.__wrapped__ = func2
+        with self.assertRaisesRegex(ValueError, 'wrapper loop'):
+            inspect.unwrap(func1)
+        with self.assertRaisesRegex(ValueError, 'wrapper loop'):
+            inspect.unwrap(func2)
+
+    def test_unhashable(self):
+        def func(): pass
+        func.__wrapped__ = None
+        class C:
+            __hash__ = None
+            __wrapped__ = func
+        self.assertIsNone(inspect.unwrap(C()))
+
+
 def test_main():
     run_unittest(
         TestDecorators, TestRetrievingSourceCode, TestOneliners, TestBuggyCases,
@@ -2291,7 +2363,7 @@ def test_main():
         TestGetcallargsFunctions, TestGetcallargsMethods,
         TestGetcallargsUnboundMethods, TestGetattrStatic, TestGetGeneratorState,
         TestNoEOL, TestSignatureObject, TestSignatureBind, TestParameterObject,
-        TestBoundArguments, TestGetClosureVars
+        TestBoundArguments, TestGetClosureVars, TestUnwrap
     )
 
 if __name__ == "__main__":
diff --git a/Misc/NEWS b/Misc/NEWS
index 80bf9fd..03c9a8e 100644
--- a/Misc/NEWS
+++ b/Misc/NEWS
@@ -171,6 +171,9 @@ Core and Builtins
 Library
 -------
 
+- Issue #13266: Added inspect.unwrap to easily unravel __wrapped__ chains
+  (initial patch by Daniel Urban and Aaron Iles)
+
 - Issue #18561: Skip name in ctypes' _build_callargs() if name is NULL.
 
 - Issue #18559: Fix NULL pointer dereference error in _pickle module