mirror of https://github.com/m-labs/artiq
59 lines
1.5 KiB
Python
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")
|