107 lines
2.8 KiB
Python
107 lines
2.8 KiB
Python
|
## @package context
|
||
|
|
# Module caffe2.python.context
|
||
|
|
|
||
|
|
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
|