From 5a3bf4894fc5fb4dae36751434cbeee3bc63361f Mon Sep 17 00:00:00 2001
From: Sebastien Bourdeauducq <sb@m-labs.hk>
Date: Thu, 11 Nov 2021 20:35:06 +0800
Subject: [PATCH] reinstate import_cache hack (#416)

---
 artiq/frontend/artiq_run.py    |  2 +-
 artiq/language/__init__.py     |  3 +-
 artiq/language/core.py         |  6 +++-
 artiq/language/import_cache.py | 52 ++++++++++++++++++++++++++++++++++
 artiq/master/worker_impl.py    |  2 +-
 5 files changed, 61 insertions(+), 4 deletions(-)
 create mode 100644 artiq/language/import_cache.py

diff --git a/artiq/frontend/artiq_run.py b/artiq/frontend/artiq_run.py
index 60a85e664..45ee1f34d 100755
--- a/artiq/frontend/artiq_run.py
+++ b/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 *
 
 
diff --git a/artiq/language/__init__.py b/artiq/language/__init__.py
index fef1a93ca..99c019a75 100644
--- a/artiq/language/__init__.py
+++ b/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__)
diff --git a/artiq/language/core.py b/artiq/language/core.py
index c6a96ecbe..ba3c91083 100644
--- a/artiq/language/core.py
+++ b/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)
 
 
diff --git a/artiq/language/import_cache.py b/artiq/language/import_cache.py
new file mode 100644
index 000000000..08de378c4
--- /dev/null
+++ b/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")
diff --git a/artiq/master/worker_impl.py b/artiq/master/worker_impl.py
index 0467a8096..29ec4e427 100644
--- a/artiq/master/worker_impl.py
+++ b/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