xref: /aosp_15_r20/external/pytorch/torch/_dynamo/code_context.py (revision da0073e96a02ea20f0ac840b70461e3646d07c45) 
1# mypy: allow-untyped-defs
2import types
3
4from .utils import ExactWeakKeyDictionary
5
6
7class CodeContextDict:
8    def __init__(self) -> None:
9        self.code_context = ExactWeakKeyDictionary()
10
11    def has_context(self, code: types.CodeType):
12        return code in self.code_context
13
14    def get_context(self, code: types.CodeType):
15        ctx = self.code_context.get(code)
16        if ctx is None:
17            ctx = {}
18            self.code_context[code] = ctx
19        return ctx
20
21    def pop_context(self, code: types.CodeType):
22        ctx = self.get_context(code)
23        self.code_context._remove_id(id(code))
24        return ctx
25
26    def clear(self):
27        self.code_context.clear()
28
29
30code_context = CodeContextDict()
31