forked from M-Labs/artiq
1
0
Fork 0
artiq/artiq/compiler/import_cache.py

59 lines
1.5 KiB
Python

import sys
import builtins
import linecache
import tokenize
import logging
import importlib.machinery as im
from artiq.experiment import kernel, portable
__all__ = ["install_hook"]
logger = logging.getLogger(__name__)
cache = dict()
im_exec_module = None
linecache_getlines = None
def hook_exec_module(self, module):
im_exec_module(self, module)
if (hasattr(module, "__file__")
# Heuristic to determine if the module may contain ARTIQ kernels.
# This breaks if kernel is not imported the usual way.
and ((getattr(module, "kernel", None) is kernel)
or (getattr(module, "portable", None) is portable))):
fn = module.__file__
try:
with tokenize.open(fn) as fp:
lines = fp.readlines()
if lines and not lines[-1].endswith("\n"):
lines[-1] += "\n"
cache[fn] = lines
except:
logger.warning("failed to add '%s' to cache", fn, exc_info=True)
else:
logger.debug("added '%s' to cache", fn)
def hook_getlines(filename, module_globals=None):
if filename in cache:
return cache[filename]
else:
return linecache_getlines(filename, module_globals)
def install_hook():
global im_exec_module, linecache_getlines
im_exec_module = im.SourceFileLoader.exec_module
im.SourceFileLoader.exec_module = hook_exec_module
linecache_getlines = linecache.getlines
linecache.getlines = hook_getlines
logger.debug("hook installed")