diff options
author | George Burgess IV <gbiv@google.com> | 2019-08-15 16:23:01 -0700 |
---|---|---|
committer | George Burgess <gbiv@chromium.org> | 2019-08-16 21:42:16 +0000 |
commit | 16603d5a467634a77b551daa2452f658d65d0ec2 (patch) | |
tree | a33816cd0796c997282b83e6b153923f22442b1c /cros_utils | |
parent | e66aac08f0d52c15652188ef4eddf0cfb8690813 (diff) | |
download | toolchain-utils-16603d5a467634a77b551daa2452f658d65d0ec2.tar.gz |
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>
Diffstat (limited to 'cros_utils')
-rw-r--r-- | cros_utils/contextlib3.py | 116 | ||||
-rwxr-xr-x | cros_utils/contextlib3_test.py | 195 |
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() |