/
manager.py
47 lines (38 loc) · 1.3 KB
/
manager.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
import weakref
from numba.core import types
class DataModelManager(object):
"""Manages mapping of FE types to their corresponding data model
"""
def __init__(self):
# { numba type class -> model factory }
self._handlers = {}
# { numba type instance -> model instance }
self._cache = weakref.WeakKeyDictionary()
def register(self, fetypecls, handler):
"""Register the datamodel factory corresponding to a frontend-type class
"""
assert issubclass(fetypecls, types.Type)
self._handlers[fetypecls] = handler
def lookup(self, fetype):
"""Returns the corresponding datamodel given the frontend-type instance
"""
try:
return self._cache[fetype]
except KeyError:
pass
handler = self._handlers[type(fetype)]
model = self._cache[fetype] = handler(self, fetype)
return model
def __getitem__(self, fetype):
"""Shorthand for lookup()
"""
return self.lookup(fetype)
def copy(self):
"""
Make a copy of the manager.
Use this to inherit from the default data model and specialize it
for custom target.
"""
dmm = DataModelManager()
dmm._handlers = self._handlers.copy()
return dmm