| | |
| | |
| |
|
| | import inspect |
| | import threading |
| | import functools |
| |
|
| |
|
| | class _ContextInfo(object): |
| | def __init__(self, cls, allow_default): |
| | self.cls = cls |
| | self.allow_default = allow_default |
| | self._local_stack = threading.local() |
| |
|
| | @property |
| | def _stack(self): |
| | if not hasattr(self._local_stack, 'obj'): |
| | self._local_stack.obj = [] |
| | return self._local_stack.obj |
| |
|
| | def enter(self, value): |
| | self._stack.append(value) |
| |
|
| | def exit(self, value): |
| | assert len(self._stack) > 0, 'Context %s is empty.' % self.cls |
| | assert self._stack.pop() == value |
| |
|
| | def get_active(self, required=True): |
| | if len(self._stack) == 0: |
| | if not required: |
| | return None |
| | assert self.allow_default, ( |
| | 'Context %s is required but none is active.' % self.cls) |
| | self.enter(self.cls()) |
| | return self._stack[-1] |
| |
|
| |
|
| | class _ContextRegistry(object): |
| | def __init__(self): |
| | self._ctxs = {} |
| |
|
| | def get(self, cls): |
| | if cls not in self._ctxs: |
| | assert issubclass(cls, Managed), "must be a context managed class, got {}".format(cls) |
| | self._ctxs[cls] = _ContextInfo(cls, allow_default=issubclass(cls, DefaultManaged)) |
| | return self._ctxs[cls] |
| |
|
| |
|
| | _CONTEXT_REGISTRY = _ContextRegistry() |
| |
|
| |
|
| | def _context_registry(): |
| | global _CONTEXT_REGISTRY |
| | return _CONTEXT_REGISTRY |
| |
|
| |
|
| | def _get_managed_classes(obj): |
| | return [ |
| | cls for cls in inspect.getmro(obj.__class__) |
| | if issubclass(cls, Managed) and cls != Managed and cls != DefaultManaged |
| | ] |
| |
|
| |
|
| |
|
| | class Managed(object): |
| | """ |
| | Managed makes the inheritted class a context managed class. |
| | |
| | class Foo(Managed): ... |
| | |
| | with Foo() as f: |
| | assert f == Foo.current() |
| | """ |
| |
|
| | @classmethod |
| | def current(cls, value=None, required=True): |
| | ctx_info = _context_registry().get(cls) |
| | if value is not None: |
| | assert isinstance(value, cls), ( |
| | 'Wrong context type. Expected: %s, got %s.' % (cls, type(value))) |
| | return value |
| | return ctx_info.get_active(required=required) |
| |
|
| | def __enter__(self): |
| | for cls in _get_managed_classes(self): |
| | _context_registry().get(cls).enter(self) |
| | return self |
| |
|
| | def __exit__(self, *args): |
| | for cls in _get_managed_classes(self): |
| | _context_registry().get(cls).exit(self) |
| |
|
| | def __call__(self, func): |
| | @functools.wraps(func) |
| | def wrapper(*args, **kwargs): |
| | with self: |
| | return func(*args, **kwargs) |
| | return wrapper |
| |
|
| |
|
| | class DefaultManaged(Managed): |
| | """ |
| | DefaultManaged is similar to Managed but if there is no parent when |
| | current() is called it makes a new one. |
| | """ |
| | pass |
| |
|
