Fix the wrongheaded implementation of context management in the decimal module and add unit tests. (python-dev discussion is ongoing regarding what we do about Python 2.5)

Doc/lib/libdecimal.tex

@@ -443,18 +443,26 @@ the \function{getcontext()} and \function{setcontext()} functions:
 \end{funcdesc}  
 
 Beginning with Python 2.5, you can also use the \keyword{with} statement
-to temporarily change the active context. For example the following code
-increases the current decimal precision by 2 places, performs a
-calculation, and then automatically restores the previous context:
+to temporarily change the active context.
 
+\begin{funcdesc}{localcontext}{\optional{c}}
+  Return a context manager that will set the current context for
+  the active thread to a copy of \var{c} on entry to the with statement
+  and restore the previous context when exiting the with statement.
+
+  For example the following code increases the current decimal precision
+  by 2 places, performs a calculation, and then automatically restores
+  the previous context:
 \begin{verbatim}
     from __future__ import with_statement
     import decimal
 
-with decimal.getcontext() as ctx:
+    with decimal.localcontext() as ctx:
         ctx.prec += 2   # add 2 more digits of precision
-    calculate_something()
+        s = calculate_something()
+    s = +s  # Round the final result back to the default precision
 \end{verbatim}
+\end{funcdesc}
 
 The context that's active in the body of the \keyword{with} statement is
 a \emph{copy} of the context you provided to the \keyword{with}

Lib/decimal.py

@@ -130,8 +130,11 @@ __all__ = [
     'ROUND_DOWN', 'ROUND_HALF_UP', 'ROUND_HALF_EVEN', 'ROUND_CEILING',
     'ROUND_FLOOR', 'ROUND_UP', 'ROUND_HALF_DOWN',
 
+    # helper for context management
+    'ContextManager',
+
     # Functions for manipulating contexts
-    'setcontext', 'getcontext'
+    'setcontext', 'getcontext', 'localcontext'
 ]
 
 import copy as _copy
@@ -458,6 +461,49 @@ else:
 
     del threading, local        # Don't contaminate the namespace
 
+def localcontext(ctx=None):
+    """Return a context manager for a copy of the supplied context
+
+    Uses a copy of the current context if no context is specified
+    The returned context manager creates a local decimal context
+    in a with statement:
+        def sin(x):
+             with localcontext() as ctx:
+                 ctx.prec += 2
+                 # Rest of sin calculation algorithm
+                 # uses a precision 2 greater than normal
+             return +s # Convert result to normal precision
+
+         def sin(x):
+             with localcontext(ExtendedContext):
+                 # Rest of sin calculation algorithm
+                 # uses the Extended Context from the
+                 # General Decimal Arithmetic Specification
+             return +s # Convert result to normal context
+
+    """
+    # The below can't be included in the docstring until Python 2.6
+    # as the doctest module doesn't understand __future__ statements
+    """
+    >>> from __future__ import with_statement
+    >>> print getcontext().prec
+    28
+    >>> with localcontext():
+    ...     ctx = getcontext()
+    ...     ctx.prec() += 2
+    ...     print ctx.prec
+    ...
+    30
+    >>> with localcontext(ExtendedContext):
+    ...     print getcontext().prec
+    ...
+    9
+    >>> print getcontext().prec
+    28
+    """
+    if ctx is None: ctx = getcontext().copy()
+    return ContextManager(ctx.copy())
+
 
 ##### Decimal class ###########################################
 
@@ -2174,19 +2220,26 @@ for name in rounding_functions:
 del name, val, globalname, rounding_functions
 
 class ContextManager(object):
-    """Helper class to simplify Context management.
+    """Context manager class to support localcontext().
 
+      Sets the supplied context in __enter__() and restores
+      the previous decimal context in __exit__()
+
+    """
+    # The below can't be included in the docstring until Python 2.6
+    # as the doctest module doesn't understand __future__ statements
+    """
     Sample usage:
-
-    with decimal.ExtendedContext:
-        s = ...
-    return +s # Convert result to normal precision
-
-    with decimal.getcontext() as ctx:
-        ctx.prec += 2
-        s = ...
-    return +s
-
+    >>> from __future__ import with_statement
+    >>> print getcontext().prec
+    28
+    >>> ctx = Context(prec=15)
+    >>> with ContextManager(ctx):
+    ...     print getcontext().prec
+    ...
+    15
+    >>> print getcontext().prec
+    28
     """
     def __init__(self, new_context):
         self.new_context = new_context
@@ -2248,9 +2301,6 @@ class Context(object):
         s.append('traps=[' + ', '.join([t.__name__ for t, v in self.traps.items() if v]) + ']')
         return ', '.join(s) + ')'
 
-    def get_manager(self):
-        return ContextManager(self.copy())
-
     def clear_flags(self):
         """Reset all flags to zero"""
         for flag in self.flags:

Lib/test/test_decimal.py

@@ -23,6 +23,7 @@ or Behaviour) to test each part, or without parameter to test both parts. If
 you're working through IDLE, you can import this test module and call test_main()
 with the corresponding argument.
 """
+from __future__ import with_statement
 
 import unittest
 import glob
@@ -1064,6 +1065,43 @@ class ContextAPItests(unittest.TestCase):
         self.assertNotEqual(id(c.flags), id(d.flags))
         self.assertNotEqual(id(c.traps), id(d.traps))
 
+class WithStatementTest(unittest.TestCase):
+    # Can't do these as docstrings until Python 2.6
+    # as doctest can't handle __future__ statements
+    def test_ContextManager(self):
+        # The basic context manager uses the supplied context
+        # without making a copy of it
+        orig_ctx = getcontext()
+        new_ctx = Context()
+        with ContextManager(new_ctx) as enter_ctx:
+            set_ctx = getcontext()
+        final_ctx = getcontext()
+        self.assert_(orig_ctx is final_ctx, 'did not restore context correctly')
+        self.assert_(new_ctx is set_ctx, 'did not set correct context')
+        self.assert_(set_ctx is enter_ctx, '__enter__ returned wrong context')
+
+    def test_localcontext(self):
+        # The helper function makes a copy of the supplied context
+        orig_ctx = getcontext()
+        with localcontext() as enter_ctx:
+            set_ctx = getcontext()
+        final_ctx = getcontext()
+        self.assert_(orig_ctx is final_ctx, 'did not restore context correctly')
+        self.assert_(orig_ctx is not set_ctx, 'did not copy the context')
+        self.assert_(set_ctx is enter_ctx, '__enter__ returned wrong context')
+
+    def test_localcontextarg(self):
+        # The helper function makes a copy of the supplied context
+        orig_ctx = getcontext()
+        new_ctx = Context(prec=42)
+        with localcontext(new_ctx) as enter_ctx:
+            set_ctx = getcontext()
+        final_ctx = getcontext()
+        self.assert_(orig_ctx is final_ctx, 'did not restore context correctly')
+        self.assert_(set_ctx.prec == new_ctx.prec, 'did not set correct context')
+        self.assert_(new_ctx is not set_ctx, 'did not copy the context')
+        self.assert_(set_ctx is enter_ctx, '__enter__ returned wrong context')
+
 def test_main(arith=False, verbose=None):
     """ Execute the tests.
 
@@ -1084,6 +1122,7 @@ def test_main(arith=False, verbose=None):
         DecimalPythonAPItests,
         ContextAPItests,
         DecimalTest,
+        WithStatementTest,
     ]
 
     try: