mirror of
https://github.com/saymrwulf/pytorch.git
synced 2026-05-15 21:00:47 +00:00
dc7d8a889e
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/48340 This changes the context managed classes from using a decorator to define them to using inheritance. Inheritance allows the python static type checking to work correctly. ``` context.define_context() class Bar(object): ... context.define_context(allow_default=True) class Foo(object): ... ``` becomes ``` class Foo(context.Managed): ... class Bar(context.DefaultManaged): ... ``` Behavior differences: * arg_name has been removed since it's not used anywhere * classes need to call `super()` in `__enter__/__exit__` methods if they override (none do) This also defines a context.pyi file to add types for python3. python2 support should not be affected Test Plan: ci buck test //caffe2/caffe2/python:context_test //caffe2/caffe2/python:checkpoint_test Reviewed By: dongyuzheng Differential Revision: D25133469 fbshipit-source-id: 16368bf723eeb6ce3308d6827f5ac5e955b4e29a
106 lines
2.8 KiB
Python
106 lines
2.8 KiB
Python
## @package context
|
|
# Module caffe2.python.context
|
|
|
|
import inspect
|
|
import threading
|
|
import six
|
|
|
|
|
|
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):
|
|
@six.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
|