reinstate import_cache hack (#416)

artiq/frontend/artiq_run.py

@@ -16,7 +16,7 @@ from artiq import __version__ as artiq_version
 from artiq.language.environment import EnvExperiment, ProcessArgumentManager
 from artiq.master.databases import DeviceDB, DatasetDB
 from artiq.master.worker_db import DeviceManager, DatasetManager
-from artiq.compiler import import_cache
+from artiq.language import import_cache
 from artiq.tools import *
 
 

artiq/language/__init__.py

@@ -3,8 +3,9 @@ from artiq.language.core import *
 from artiq.language.environment import *
 from artiq.language.units import *
 from artiq.language.scan import *
+from . import import_cache
 
-__all__ = []
+__all__ = ["import_cache"]
 __all__.extend(core.__all__)
 __all__.extend(environment.__all__)
 __all__.extend(units.__all__)

artiq/language/core.py

@@ -4,9 +4,11 @@ Core ARTIQ extensions to the Python language.
 
 from typing import Generic, TypeVar
 from functools import wraps
-from inspect import getfullargspec
+from inspect import getfullargspec, getmodule
 from types import SimpleNamespace
 
+from artiq.language import import_cache
+
 
 __all__ = [
     "KernelInvariant", "round64",
@@ -32,10 +34,12 @@ _registered_classes = set()
 
 def _register_function(fun):
     assert _allow_registration
+    import_cache.add_module_to_cache(getmodule(fun))
     _registered_functions.add(fun)
 
 def _register_class(cls):
     assert _allow_registration
+    import_cache.add_module_to_cache(getmodule(cls))
     _registered_classes.add(cls)
 
 

artiq/language/import_cache.py

@@ -0,0 +1,52 @@
+"""
+Caches source files on import so that inspect.getsource returns the source code that
+was imported (or at least with a small race window), not what is currently on the
+filesystem at the time inspect.getsource is called.
+This is a hack and it still has races, and it would be better if Python supported
+this correctly, but it does not.
+"""
+
+import linecache
+import tokenize
+import logging
+
+
+__all__ = ["install_hook", "add_module_to_cache"]
+
+
+logger = logging.getLogger(__name__)
+
+
+cache = dict()
+linecache_getlines = None
+
+
+def add_module_to_cache(module):
+    if hasattr(module, "__file__"):
+        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 linecache_getlines
+
+    linecache_getlines = linecache.getlines
+    linecache.getlines = hook_getlines
+
+    logger.debug("hook installed")

artiq/master/worker_impl.py

@@ -28,7 +28,7 @@ from artiq.language.environment import (
 )
 from artiq.language.core import set_watchdog_factory, TerminationRequested
 from artiq.language.types import TBool
-from artiq.compiler import import_cache
+from artiq.language import import_cache
 from artiq.coredevice.core import CompileError, host_only, _render_diagnostic
 from artiq import __version__ as artiq_version