Add functools.update_wrapper() and functools.wraps() as described in PEP 356
diff --git a/Lib/functools.py b/Lib/functools.py
index 4935c9f..8783f08 100644
--- a/Lib/functools.py
+++ b/Lib/functools.py
@@ -1,26 +1,51 @@
-"""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)
diff --git a/Lib/test/test_functools.py b/Lib/test/test_functools.py
index 609e8f4..8dc185b 100644
--- a/Lib/test/test_functools.py
+++ b/Lib/test/test_functools.py
@@ -152,6 +152,113 @@
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 @@
TestPartial,
TestPartialSubclass,
TestPythonPartial,
+ TestUpdateWrapper,
+ TestWraps
)
test_support.run_unittest(*test_classes)