From 911b459522d9dee379cdc64a87bbea2f22fdfbf9 Mon Sep 17 00:00:00 2001
From: Nick Coghlan <ncoghlan@gmail.com>
Date: Mon, 2 May 2016 16:41:56 +1000
Subject: [PATCH] Issue #5: support old-style classes in ExitStack

---
 NEWS.rst            |  4 +++-
 contextlib2.py      | 18 ++++++++++++++++--
 docs/index.rst      |  4 ----
 test_contextlib2.py | 33 +++++++++++++++++++++++++++++++++
 4 files changed, 52 insertions(+), 7 deletions(-)

diff --git a/NEWS.rst b/NEWS.rst
index 6b43320..f8451e0 100644
--- a/NEWS.rst
+++ b/NEWS.rst
@@ -4,7 +4,9 @@ Release History
 0.5.3 (not yet released)
 ^^^^^^^^^^^^^^^^^^^^^^^^
 
-* TBD
+* ``ExitStack`` now correctly handles context managers implemented as old-style
+  classes in Python 2.x (such as ``codecs.StreamReader`` and
+  ``codecs.StreamWriter``)
 
 0.5.2 (2016-05-02)
 ^^^^^^^^^^^^^^^^^^
diff --git a/contextlib2.py b/contextlib2.py
index 133b8e4..a6acf65 100644
--- a/contextlib2.py
+++ b/contextlib2.py
@@ -288,6 +288,20 @@ else:
         exc_type, exc_value, exc_tb = exc_details
         exec ("raise exc_type, exc_value, exc_tb")
 
+# Handle old-style classes if they exist
+try:
+    from types import InstanceType
+except ImportError:
+    # Python 3 doesn't have old-style classes
+    _get_type = type
+else:
+    # Need to handle old-style context managers on Python 2
+    def _get_type(obj):
+        obj_type = type(obj)
+        if obj_type is InstanceType:
+            return obj.__class__ # Old-style class
+        return obj_type # New-style class
+
 # Inspired by discussions on http://bugs.python.org/issue13585
 class ExitStack(object):
     """Context manager for dynamic management of a stack of exit callbacks
@@ -328,7 +342,7 @@ class ExitStack(object):
         """
         # We use an unbound method rather than a bound method to follow
         # the standard lookup behaviour for special methods
-        _cb_type = type(exit)
+        _cb_type = _get_type(exit)
         try:
             exit_method = _cb_type.__exit__
         except AttributeError:
@@ -358,7 +372,7 @@ class ExitStack(object):
         returns the result of the __enter__ method.
         """
         # We look up the special methods on the type to match the with statement
-        _cm_type = type(cm)
+        _cm_type = _get_type(cm)
         _exit = _cm_type.__exit__
         result = _cm_type.__enter__(cm)
         self._push_cm_exit(cm, _exit)
diff --git a/docs/index.rst b/docs/index.rst
index b25288b..71474a2 100644
--- a/docs/index.rst
+++ b/docs/index.rst
@@ -310,10 +310,6 @@ Functions and classes provided:
    foundation for higher level context managers that manipulate the exit
    stack in application specific ways.
 
-   Context managers used with :class:`ExitStack` must be new-style classes -
-   this is the default on Python 3, but requires explicitly inheriting from
-   :class:`object` or another new-style class in Python 2.
-
    .. versionadded:: 0.4
       Part of the standard library in Python 3.3 and later
 
diff --git a/test_contextlib2.py b/test_contextlib2.py
index 478b3aa..3aba52f 100755
--- a/test_contextlib2.py
+++ b/test_contextlib2.py
@@ -706,6 +706,39 @@ class TestExitStack(unittest.TestCase):
         stack.push(cm)
         self.assertIs(stack._exit_callbacks[-1], cm)
 
+    def test_default_class_semantics(self):
+        # For Python 2.x, this ensures compatibility with old-style classes
+        # For Python 3.x, it just reruns some of the other tests
+        class DefaultCM:
+            def __enter__(self):
+                result.append("Enter")
+            def __exit__(self, *exc_details):
+                result.append("Exit")
+        class DefaultCallable:
+            def __call__(self, *exc_details):
+                result.append("Callback")
+
+        result = []
+        cm = DefaultCM()
+        cb = DefaultCallable()
+        with ExitStack() as stack:
+            stack.enter_context(cm)
+            self.assertIs(stack._exit_callbacks[-1].__self__, cm)
+            stack.push(cb)
+            stack.push(cm)
+            self.assertIs(stack._exit_callbacks[-1].__self__, cm)
+            result.append("Running")
+            stack.callback(cb)
+            self.assertIs(stack._exit_callbacks[-1].__wrapped__, cb)
+        self.assertEqual(result, ["Enter", "Running",
+                                  "Callback", "Exit",
+                                  "Callback", "Exit",
+                                 ])
+
+        with ExitStack():
+            pass
+
+
 
 class TestRedirectStream: