Note that we can make it a mixin class too for easier usage:
from abc import ABC, abstractmethod
from contextlib import contextmanager
class ContextManagerMixin(ABC):
def __init__(self):
self._contexts = []
def __init_subclass__(cls):
cls.__with__ = contextmanager(cls.__with__)
@abstractmethod
def __with__(self): ...
def __enter__(self):
context = self.__with__()
self._contexts.append(context)
return context.__enter__()
def __exit__(self, exc_type, exc_value, traceback):
return self._contexts.pop().__exit__(exc_type, exc_value, traceback)
so that:
class Counter(ContextManagerMixin):
def __init__(self):
super().__init__()
self.count = 0
def __with__(self):
self.count += 1
yield self
self.count -= 1
counter = Counter()
with counter as counter1:
with counter as counter2:
print(counter2.count) # outputs 2
print(counter1.count) # outputs 1