Add functools.update_wrapper() and functools.wraps() as described in PEP 356

Doc/lib/libfunctools.tex

@@ -5,6 +5,7 @@
 
 \moduleauthor{Peter Harris}{scav@blueyonder.co.uk}
 \moduleauthor{Raymond Hettinger}{python@rcn.com}
+\moduleauthor{Nick Coghlan}{ncoghlan@gmail.com}
 \sectionauthor{Peter Harris}{scav@blueyonder.co.uk}
 
 \modulesynopsis{Higher-order functions and operations on callable objects.}
@@ -50,6 +51,51 @@ two:
   \end{verbatim}
 \end{funcdesc}
 
+\begin{funcdesc}{update_wrapper}
+{wrapper, wrapped\optional{, assigned}\optional{, updated}}
+Update a wrapper function to look like the wrapped function. The optional
+arguments are tuples to specify which attributes of the original
+function are assigned directly to the matching attributes on the wrapper
+function and which attributes of the wrapper function are updated with
+the corresponding attributes from the original function. The default
+values for these arguments are the module level constants
+\var{WRAPPER_ASSIGNMENTS} (which assigns to the wrapper function's name,
+module and documentation string) and \var{WRAPPER_UPDATES} (which
+updates the wrapper function's instance dictionary).
+
+The main intended use for this function is in decorator functions
+which wrap the decorated function and return the wrapper. If the
+wrapper function is not updated, the metadata of the returned function
+will reflect the wrapper definition rather than the original function
+definition, which is typically less than helpful.
+\end{funcdesc}
+
+\begin{funcdesc}{wraps}
+{wrapped\optional{, assigned}\optional{, updated}}
+This is a convenience function for invoking
+\code{partial(update_wrapper, wrapped=wrapped, assigned=assigned, updated=updated)}
+as a function decorator when defining a wrapper function. For example:
+  \begin{verbatim}
+        >>> def my_decorator(f):
+        ...     @wraps(f)
+        ...     def wrapper(*args, **kwds):
+        ...         print 'Calling decorated function'
+        ...         return f(*args, **kwds)
+        ...     return wrapper
+        ...
+        >>> @my_decorator
+        ... def example():
+        ...     print 'Called example function'
+        ...
+        >>> example()
+        Calling decorated function
+        Called example function
+        >>> example.__name__
+        'example'
+  \end{verbatim}
+Without the use of this decorator factory, the name of the example
+function would have been \code{'wrapper'}.
+\end{funcdesc}
 
 
 \subsection{\class{partial} Objects \label{partial-objects}}

Lib/functools.py

-"""functools.py - Tools for working with functions
+"""functools.py - Tools for working with functions and callable objects
 """
 # Python module wrapper for _functools C module
 # to allow utilities written in Python to be added
 # to the functools module.
 # Written by Nick Coghlan <ncoghlan at gmail.com>
-#   Copyright (c) 2006 Python Software Foundation.
+#   Copyright (C) 2006 Python Software Foundation.
+# See C source code for _functools credits/copyright
 
 from _functools import partial
-__all__ = [
-    "partial",
-]
 
-# Still to come here (need to write tests and docs):
-#   update_wrapper - utility function to transfer basic function
-#                    metadata to wrapper functions
-#   WRAPPER_ASSIGNMENTS & WRAPPER_UPDATES - defaults args to above
-#           (update_wrapper has been approved by BDFL)
-#   wraps - decorator factory equivalent to:
-#               def wraps(f):
-#                     return partial(update_wrapper, wrapped=f)
-#
-# The wraps function makes it easy to avoid the bug that afflicts the
-# decorator example in the python-dev email proposing the
-# update_wrapper function:
-# http://mail.python.org/pipermail/python-dev/2006-May/064775.html
+# update_wrapper() and wraps() are tools to help write
+# wrapper functions that can handle naive introspection
+
+WRAPPER_ASSIGNMENTS = ('__module__', '__name__', '__doc__')
+WRAPPER_UPDATES = ('__dict__',)
+def update_wrapper(wrapper,
+                   wrapped,
+                   assigned = WRAPPER_ASSIGNMENTS,
+                   updated = WRAPPER_UPDATES):
+    """Update a wrapper function to look like the wrapped function
+
+       wrapper is the function to be updated
+       wrapped is the original function
+       assigned is a tuple naming the attributes assigned directly
+       from the wrapped function to the wrapper function (defaults to
+       functools.WRAPPER_ASSIGNMENTS)
+       updated is a tuple naming the attributes off the wrapper that
+       are updated with the corresponding attribute from the wrapped
+       function (defaults to functools.WRAPPER_UPDATES)
+    """
+    for attr in assigned:
+        setattr(wrapper, attr, getattr(wrapped, attr))
+    for attr in updated:
+        getattr(wrapper, attr).update(getattr(wrapped, attr))
+    # Return the wrapper so this can be used as a decorator via partial()
+    return wrapper
+
+def wraps(wrapped,
+          assigned = WRAPPER_ASSIGNMENTS,
+          updated = WRAPPER_UPDATES):
+    """Decorator factory to apply update_wrapper() to a wrapper function
+
+       Returns a decorator that invokes update_wrapper() with the decorated
+       function as the wrapper argument and the arguments to wraps() as the
+       remaining arguments. Default arguments are as for update_wrapper().
+       This is a convenience function to simplify applying partial() to
+       update_wrapper().
+    """
+    return partial(update_wrapper, wrapped=wrapped,
+                   assigned=assigned, updated=updated)

Lib/test/test_functools.py

@@ -152,6 +152,113 @@ class TestPythonPartial(TestPartial):
 
     thetype = PythonPartial
 
+class TestUpdateWrapper(unittest.TestCase):
+
+    def check_wrapper(self, wrapper, wrapped,
+                      assigned=functools.WRAPPER_ASSIGNMENTS,
+                      updated=functools.WRAPPER_UPDATES):
+        # Check attributes were assigned
+        for name in assigned:
+            self.failUnless(getattr(wrapper, name) is getattr(wrapped, name))
+        # Check attributes were updated
+        for name in updated:
+            wrapper_attr = getattr(wrapper, name)
+            wrapped_attr = getattr(wrapped, name)
+            for key in wrapped_attr:
+                self.failUnless(wrapped_attr[key] is wrapper_attr[key])
+
+    def test_default_update(self):
+        def f():
+            """This is a test"""
+            pass
+        f.attr = 'This is also a test'
+        def wrapper():
+            pass
+        functools.update_wrapper(wrapper, f)
+        self.check_wrapper(wrapper, f)
+        self.assertEqual(wrapper.__name__, 'f')
+        self.assertEqual(wrapper.__doc__, 'This is a test')
+        self.assertEqual(wrapper.attr, 'This is also a test')
+
+    def test_no_update(self):
+        def f():
+            """This is a test"""
+            pass
+        f.attr = 'This is also a test'
+        def wrapper():
+            pass
+        functools.update_wrapper(wrapper, f, (), ())
+        self.check_wrapper(wrapper, f, (), ())
+        self.assertEqual(wrapper.__name__, 'wrapper')
+        self.assertEqual(wrapper.__doc__, None)
+        self.failIf(hasattr(wrapper, 'attr'))
+
+    def test_selective_update(self):
+        def f():
+            pass
+        f.attr = 'This is a different test'
+        f.dict_attr = dict(a=1, b=2, c=3)
+        def wrapper():
+            pass
+        wrapper.dict_attr = {}
+        assign = ('attr',)
+        update = ('dict_attr',)
+        functools.update_wrapper(wrapper, f, assign, update)
+        self.check_wrapper(wrapper, f, assign, update)
+        self.assertEqual(wrapper.__name__, 'wrapper')
+        self.assertEqual(wrapper.__doc__, None)
+        self.assertEqual(wrapper.attr, 'This is a different test')
+        self.assertEqual(wrapper.dict_attr, f.dict_attr)
+
+
+class TestWraps(TestUpdateWrapper):
+
+    def test_default_update(self):
+        def f():
+            """This is a test"""
+            pass
+        f.attr = 'This is also a test'
+        @functools.wraps(f)
+        def wrapper():
+            pass
+        self.check_wrapper(wrapper, f)
+        self.assertEqual(wrapper.__name__, 'f')
+        self.assertEqual(wrapper.__doc__, 'This is a test')
+        self.assertEqual(wrapper.attr, 'This is also a test')
+
+    def test_no_update(self):
+        def f():
+            """This is a test"""
+            pass
+        f.attr = 'This is also a test'
+        @functools.wraps(f, (), ())
+        def wrapper():
+            pass
+        self.check_wrapper(wrapper, f, (), ())
+        self.assertEqual(wrapper.__name__, 'wrapper')
+        self.assertEqual(wrapper.__doc__, None)
+        self.failIf(hasattr(wrapper, 'attr'))
+
+    def test_selective_update(self):
+        def f():
+            pass
+        f.attr = 'This is a different test'
+        f.dict_attr = dict(a=1, b=2, c=3)
+        def add_dict_attr(f):
+            f.dict_attr = {}
+            return f
+        assign = ('attr',)
+        update = ('dict_attr',)
+        @functools.wraps(f, assign, update)
+        @add_dict_attr
+        def wrapper():
+            pass
+        self.check_wrapper(wrapper, f, assign, update)
+        self.assertEqual(wrapper.__name__, 'wrapper')
+        self.assertEqual(wrapper.__doc__, None)
+        self.assertEqual(wrapper.attr, 'This is a different test')
+        self.assertEqual(wrapper.dict_attr, f.dict_attr)
+
 
 
 def test_main(verbose=None):
@@ -160,6 +267,8 @@ def test_main(verbose=None):
         TestPartial,
         TestPartialSubclass,
         TestPythonPartial,
+        TestUpdateWrapper,
+        TestWraps
     )
     test_support.run_unittest(*test_classes)
 

Misc/NEWS

@@ -127,6 +127,10 @@ Extension Modules
 Library
 -------
 
+- The functions update_wrapper() and wraps() have been added to the functools
+  module. These make it easier to copy relevant metadata from the original
+  function when writing wrapper functions.
+
 - The optional ``isprivate`` argument to ``doctest.testmod()``, and the
   ``doctest.is_private()`` function, both deprecated in 2.4, were removed.