Python context manager protocol extension

cmv2.py

from __future__ import annotations


def CMProtocolV2(typ):
    # Copy from contextlib._GeneratorContextManager
    def __enter__(self):
        self._gen = self.__with__()
        try:
            return next(self._gen)
        except StopIteration:
            raise RuntimeError("generator didn't yield") from None

    def __exit__(self, typ, value, traceback):
        if typ is None:
            try:
                next(self._gen)
            except StopIteration:
                return False
            else:
                raise RuntimeError("generator didn't stop")
        else:
            if value is None:
                value = typ()
            try:
                self._gen.throw(value)
            except StopIteration as exc:
                return exc is not value
            except RuntimeError as exc:
                if exc is value:
                    exc.__traceback__ = traceback
                    return False
                if (
                    isinstance(value, StopIteration)
                    and exc.__cause__ is value
                ):
                    value.__traceback__ = traceback
                    return False
                raise
            except BaseException as exc:
                if exc is not value:
                    raise
                exc.__traceback__ = traceback
                return False
            raise RuntimeError("generator didn't stop after throw()")

    typ.__enter__ = __enter__
    typ.__exit__ = __exit__
    return typ


@CMProtocolV2
class Manager:
    def __with__(self):
        try:
            print("enter")
            yield
        except ZeroDivisionError as e:
            print("handle zero division")
            import traceback
            traceback.print_exception(e)
        else:
            print("no exception")


with Manager() as value:
    pass

with Manager() as value:
    1 / 0

with Manager() as value:
    raise Exception("not handled")