pytorch

1
## @package context
2
# Module caffe2.python.context
3

4
import inspect
5
import threading
6
import functools
7

8

9
class _ContextInfo:
10
    def __init__(self, cls, allow_default):
11
        self.cls = cls
12
        self.allow_default = allow_default
13
        self._local_stack = threading.local()
14

15
    @property
16
    def _stack(self):
17
        if not hasattr(self._local_stack, 'obj'):
18
            self._local_stack.obj = []
19
        return self._local_stack.obj
20

21
    def enter(self, value):
22
        self._stack.append(value)
23

24
    def exit(self, value):
25
        assert len(self._stack) > 0, 'Context %s is empty.' % self.cls
26
        assert self._stack.pop() == value
27

28
    def get_active(self, required=True):
29
        if len(self._stack) == 0:
30
            if not required:
31
                return None
32
            assert self.allow_default, (
33
                'Context %s is required but none is active.' % self.cls)
34
            self.enter(self.cls())
35
        return self._stack[-1]
36

37

38
class _ContextRegistry:
39
    def __init__(self):
40
        self._ctxs = {}
41

42
    def get(self, cls):
43
        if cls not in self._ctxs:
44
            assert issubclass(cls, Managed), "must be a context managed class, got {}".format(cls)
45
            self._ctxs[cls] = _ContextInfo(cls, allow_default=issubclass(cls, DefaultManaged))
46
        return self._ctxs[cls]
47

48

49
_CONTEXT_REGISTRY = _ContextRegistry()
50

51

52
def _context_registry():
53
    global _CONTEXT_REGISTRY
54
    return _CONTEXT_REGISTRY
55

56

57
def _get_managed_classes(obj):
58
    return [
59
        cls for cls in inspect.getmro(obj.__class__)
60
        if issubclass(cls, Managed) and cls != Managed and cls != DefaultManaged
61
    ]
62

63

64

65
class Managed:
66
    """
67
    Managed makes the inheritted class a context managed class.
68

69
        class Foo(Managed): ...
70

71
        with Foo() as f:
72
            assert f == Foo.current()
73
    """
74

75
    @classmethod
76
    def current(cls, value=None, required=True):
77
        ctx_info = _context_registry().get(cls)
78
        if value is not None:
79
            assert isinstance(value, cls), (
80
                'Wrong context type. Expected: %s, got %s.' % (cls, type(value)))
81
            return value
82
        return ctx_info.get_active(required=required)
83

84
    def __enter__(self):
85
        for cls in _get_managed_classes(self):
86
            _context_registry().get(cls).enter(self)
87
        return self
88

89
    def __exit__(self, *args):
90
        for cls in _get_managed_classes(self):
91
            _context_registry().get(cls).exit(self)
92

93
    def __call__(self, func):
94
        @functools.wraps(func)
95
        def wrapper(*args, **kwargs):
96
            with self:
97
                return func(*args, **kwargs)
98
        return wrapper
99

100

101
class DefaultManaged(Managed):
102
    """
103
    DefaultManaged is similar to Managed but if there is no parent when
104
    current() is called it makes a new one.
105
    """
106
    pass
107