cros_utils: add an ExitStack utility

Python3 has a really handy tool called ExitStack. It's meant to make
context managers a bit more dynamic in what they can do. Sadly, Python2
doesn't have this.

This CL adds a simple implementation of ExitStack to our utilities. We
want this feature now, and the hope is that when we move to Python3,
we'll be able to just mass-replace contextlib3 with contextlib, and
everything will Just Work.

BUG=None
TEST=Unittest

Change-Id: I1e20271da598c86c00531821e1053bc33ecc2f63
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/third_party/toolchain-utils/+/1757218
Reviewed-by: Manoj Gupta <manojgupta@chromium.org>
Tested-by: George Burgess <gbiv@chromium.org>

2 files changed, 311 insertions, 0 deletions

diff --git a/cros_utils/contextlib3.py b/cros_utils/contextlib3.py
new file mode 100644
index 00000000..9fabbf6e
--- /dev/null
+++ b/cros_utils/contextlib3.py
@@ -0,0 +1,116 @@
+# -*- coding: utf-8 -*-
+# Copyright 2019 The Chromium OS Authors. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+"""Random utilties from Python3's contextlib."""
+
+from __future__ import division
+from __future__ import print_function
+
+import sys
+
+
+class ExitStack(object):
+  """https://docs.python.org/3/library/contextlib.html#contextlib.ExitStack"""
+
+  def __init__(self):
+    self._stack = []
+    self._is_entered = False
+
+  def _assert_is_entered(self):
+    # Strictly, entering has no effect on the operations that call this.
+    # However, if you're trying to e.g. push things to an ExitStack that hasn't
+    # yet been entered, that's likely a bug.
+    assert self._is_entered, 'ExitStack op performed before entering'
+
+  def __enter__(self):
+    self._is_entered = True
+    return self
+
+  def _perform_exit(self, exc_type, exc, exc_traceback):
+    # I suppose a better name for this is
+    # `take_exception_handling_into_our_own_hands`, but that's harder to type.
+    exception_handled = False
+    while self._stack:
+      fn = self._stack.pop()
+      # The except clause below is meant to run as-if it's a `finally` block,
+      # but `finally` blocks don't have easy access to exceptions currently in
+      # flight. Hence, we do need to catch things like KeyboardInterrupt,
+      # SystemExit, ...
+      # pylint: disable=bare-except
+      try:
+        # If an __exit__ handler returns a truthy value, we should assume that
+        # it handled the exception appropriately. Otherwise, we need to keep it
+        # with us. (PEP 343)
+        if fn(exc_type, exc, exc_traceback):
+          exc_type, exc, exc_traceback = None, None, None
+          exception_handled = True
+      except:
+        # Python2 doesn't appear to have the notion of 'exception causes',
+        # which is super unfortunate. In the case:
+        #
+        # @contextlib.contextmanager
+        # def foo()
+        #   try:
+        #     yield
+        #   finally:
+        #     raise ValueError
+        #
+        # with foo():
+        #   assert False
+        #
+        # ...Python will only note the ValueError; nothing about the failing
+        # assertion is printed.
+        #
+        # I guess on the bright side, that means we don't have to fiddle with
+        # __cause__s/etc.
+        exc_type, exc, exc_traceback = sys.exc_info()
+        exception_handled = True
+
+    if not exception_handled:
+      return False
+
+    # Something changed. We either need to raise for ourselves, or note that
+    # the exception has been suppressed.
+    if exc_type is not None:
+      raise exc_type, exc, exc_traceback
+
+    # Otherwise, the exception was suppressed. Go us!
+    return True
+
+  def __exit__(self, exc_type, exc, exc_traceback):
+    return self._perform_exit(exc_type, exc, exc_traceback)
+
+  def close(self):
+    """Unwinds the exit stack, unregistering all events"""
+    self._perform_exit(None, None, None)
+
+  def enter_context(self, cm):
+    """Enters the given context manager, and registers it to be exited."""
+    self._assert_is_entered()
+
+    # The spec specifically notes that we should take __exit__ prior to calling
+    # __enter__.
+    exit_cleanup = cm.__exit__
+    result = cm.__enter__()
+    self._stack.append(exit_cleanup)
+    return result
+
+  # pylint complains about `exit` being redefined. `exit` is the documented
+  # name of this param, and renaming it would break portability if someone
+  # decided to `push(exit=foo)`, so just ignore the lint.
+  # pylint: disable=redefined-builtin
+  def push(self, exit):
+    """Like `enter_context`, but won't enter the value given."""
+    self._assert_is_entered()
+    self._stack.append(exit.__exit__)
+
+  def callback(self, callback, *args, **kwargs):
+    """Performs the given callback on exit"""
+    self._assert_is_entered()
+
+    def fn(_exc_type, _exc, _exc_traceback):
+      callback(*args, **kwargs)
+
+    self._stack.append(fn)
diff --git a/cros_utils/contextlib3_test.py b/cros_utils/contextlib3_test.py
new file mode 100755
index 00000000..76c010f2
--- /dev/null
+++ b/cros_utils/contextlib3_test.py
@@ -0,0 +1,195 @@
+#!/usr/bin/env python2
+# -*- coding: utf-8 -*-
+# Copyright 2019 The Chromium OS Authors. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+"""Tests for contextlib3"""
+
+from __future__ import division
+from __future__ import print_function
+
+import contextlib
+import unittest
+
+import contextlib3
+
+
+class SomeException(Exception):
+  """Just an alternative to ValueError in the Exception class hierarchy."""
+  pass
+
+
+class TestExitStack(unittest.TestCase):
+  """Tests contextlib3.ExitStack"""
+
+  def test_exceptions_in_exit_override_exceptions_in_with(self):
+
+    @contextlib.contextmanager
+    def raise_exit():
+      raised = False
+      try:
+        yield
+      except Exception:
+        raised = True
+        raise ValueError
+      finally:
+        self.assertTrue(raised)
+
+    # (As noted in comments in contextlib3, this behavior is consistent with
+    # how python2 works. Namely, if __exit__ raises, the exception from
+    # __exit__ overrides the inner exception)
+    with self.assertRaises(ValueError):
+      with contextlib3.ExitStack() as stack:
+        stack.enter_context(raise_exit())
+        raise SomeException()
+
+  def test_raising_in_exit_doesnt_block_later_exits(self):
+    exited = []
+
+    @contextlib.contextmanager
+    def raise_exit():
+      try:
+        yield
+      finally:
+        exited.append('raise')
+        raise ValueError
+
+    @contextlib.contextmanager
+    def push_exit():
+      try:
+        yield
+      finally:
+        exited.append('push')
+
+    with self.assertRaises(ValueError):
+      with contextlib3.ExitStack() as stack:
+        stack.enter_context(push_exit())
+        stack.enter_context(raise_exit())
+    self.assertEqual(exited, ['raise', 'push'])
+
+    exited = []
+    with self.assertRaises(ValueError):
+      with contextlib3.ExitStack() as stack:
+        stack.enter_context(push_exit())
+        stack.enter_context(raise_exit())
+        raise SomeException()
+    self.assertEqual(exited, ['raise', 'push'])
+
+  def test_push_doesnt_enter_the_context(self):
+    exited = []
+
+    test_self = self
+
+    class Manager(object):
+      """A simple ContextManager for testing purposes"""
+
+      def __enter__(self):
+        test_self.fail('context manager was entered :(')
+
+      def __exit__(self, *args, **kwargs):
+        exited.append(1)
+
+    with contextlib3.ExitStack() as stack:
+      stack.push(Manager())
+      self.assertEqual(exited, [])
+    self.assertEqual(exited, [1])
+
+  def test_callbacks_are_run_properly(self):
+    callback_was_run = []
+
+    def callback(arg, some_kwarg=None):
+      self.assertEqual(arg, 41)
+      self.assertEqual(some_kwarg, 42)
+      callback_was_run.append(1)
+
+    with contextlib3.ExitStack() as stack:
+      stack.callback(callback, 41, some_kwarg=42)
+      self.assertEqual(callback_was_run, [])
+    self.assertEqual(callback_was_run, [1])
+
+    callback_was_run = []
+    with self.assertRaises(ValueError):
+      with contextlib3.ExitStack() as stack:
+        stack.callback(callback, 41, some_kwarg=42)
+        raise ValueError()
+    self.assertEqual(callback_was_run, [1])
+
+  def test_finallys_are_run(self):
+    finally_run = []
+
+    @contextlib.contextmanager
+    def append_on_exit():
+      try:
+        yield
+      finally:
+        finally_run.append(0)
+
+    with self.assertRaises(ValueError):
+      with contextlib3.ExitStack() as stack:
+        stack.enter_context(append_on_exit())
+        raise ValueError()
+    self.assertEqual(finally_run, [0])
+
+  def test_unwinding_happens_in_reverse_order(self):
+    exit_runs = []
+
+    @contextlib.contextmanager
+    def append_things(start_push, end_push):
+      exit_runs.append(start_push)
+      try:
+        yield
+      finally:
+        exit_runs.append(end_push)
+
+    with contextlib3.ExitStack() as stack:
+      stack.enter_context(append_things(1, 4))
+      stack.enter_context(append_things(2, 3))
+    self.assertEqual(exit_runs, [1, 2, 3, 4])
+
+    exit_runs = []
+    with self.assertRaises(ValueError):
+      with contextlib3.ExitStack() as stack:
+        stack.enter_context(append_things(1, 4))
+        stack.enter_context(append_things(2, 3))
+        raise ValueError
+    self.assertEqual(exit_runs, [1, 2, 3, 4])
+
+  def test_exceptions_are_propagated(self):
+
+    @contextlib.contextmanager
+    def die_on_regular_exit():
+      yield
+      self.fail('Unreachable in theory')
+
+    with self.assertRaises(ValueError):
+      with contextlib3.ExitStack() as stack:
+        stack.enter_context(die_on_regular_exit())
+        raise ValueError()
+
+  def test_exceptions_can_be_blocked(self):
+
+    @contextlib.contextmanager
+    def block():
+      try:
+        yield
+      except Exception:
+        pass
+
+    with contextlib3.ExitStack() as stack:
+      stack.enter_context(block())
+      raise ValueError()
+
+  def test_objects_are_returned_from_enter_context(self):
+
+    @contextlib.contextmanager
+    def yield_arg(arg):
+      yield arg
+
+    with contextlib3.ExitStack() as stack:
+      val = stack.enter_context(yield_arg(1))
+      self.assertEqual(val, 1)
+
+
+if __name__ == '__main__':
+  unittest.main()