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)

1 files changed, 65 insertions, 15 deletions

diff --git a/Lib/decimal.py b/Lib/decimal.py
index 396c413..a5176e6 100644
--- a/Lib/decimal.py
+++ b/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,20 +2220,27 @@ 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().
 
-    Sample usage:
-
-    with decimal.ExtendedContext:
-        s = ...
-    return +s # Convert result to normal precision
-
-    with decimal.getcontext() as ctx:
-        ctx.prec += 2
-        s = ...
-    return +s
+      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:
+    >>> 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
     def __enter__(self):
@@ -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: