safe_contextmanager

safe_contextlib.py

import sys
from functools import wraps
from contextlib import ContextDecorator, AbstractContextManager

#  Copied and changed from python3.6/contextlib.py


def safe_contextmanager(func):
	@wraps(func)
	def helper(*args, **kwds):
		return _SafeGeneratorContextManager(func, args, kwds)

	return helper


class _SafeGeneratorContextManager(ContextDecorator, AbstractContextManager):
	"""Helper for @safe_contextmanager decorator."""

	def __init__(self, func, args, kwds):
		self.gen = func(*args, **kwds)
		self.func, self.args, self.kwds = func, args, kwds
		# Issue 19330: ensure context manager instances have good docstrings
		doc = getattr(func, "__doc__", None)
		if doc is None:
			doc = type(self).__doc__
		self.__doc__ = doc

	# Unfortunately, this still doesn't provide good help output when
	# inspecting the created context manager instances, since pydoc
	# currently bypasses the instance docstring and shows the docstring
	# for the class instead.
	# See http://bugs.python.org/issue19404 for more details.

	def _recreate_cm(self):
		# _GCM instances are one-shot context managers, so the
		# CM must be recreated each time a decorated function is
		# called
		return self.__class__(self.func, self.args, self.kwds)

	def __enter__(self):
		try:
			return next(self.gen)
		except StopIteration:
			raise RuntimeError("generator didn't yield") from None

	def __exit__(self, type, value, traceback):
		if type is None:
			try:
				next(self.gen)
			except StopIteration:
				return False
			else:
				raise RuntimeError("generator didn't stop")
		else:
			if value is None:
				# Need to force instantiation so we can reliably
				# tell if we get the same exception back
				value = type()
			try:
				next(self.gen)
			except:
				pass
			try:
				self.gen.throw(type, value, traceback)
			except StopIteration as exc:
				# Suppress StopIteration *unless* it's the same exception that
				# was passed to throw().  This prevents a StopIteration
				# raised inside the "with" statement from being suppressed.
				return exc is not value
			except RuntimeError as exc:
				# Don't re-raise the passed in exception. (issue27122)
				if exc is value:
					return False
				# Likewise, avoid suppressing if a StopIteration exception
				# was passed to throw() and later wrapped into a RuntimeError
				# (see PEP 479).
				if type is StopIteration and exc.__cause__ is value:
					return False
				raise
			except:
				# only re-raise if it's *not* the exception that was
				# passed to throw(), because __exit__() must not raise
				# an exception unless __exit__() itself failed.  But throw()
				# has to raise the exception to signal propagation, so this
				# fixes the impedance mismatch between the throw() protocol
				# and the __exit__() protocol.
				#
				if sys.exc_info()[1] is value:
					return False
				raise
			raise RuntimeError("generator didn't stop after throw()")