diff --git a/nac3artiq/demo/demo.py b/nac3artiq/demo/demo.py index 2a4b84907..75c328e95 100644 --- a/nac3artiq/demo/demo.py +++ b/nac3artiq/demo/demo.py @@ -12,7 +12,7 @@ class Demo: self.led1 = TTLOut(self.core, 19) @kernel - def run_k(self): + def run(self): self.core.reset() while True: with parallel: @@ -20,9 +20,6 @@ class Demo: self.led1.pulse(100.*ms) self.core.delay(100.*ms) - def run(self): - self.core.run(self.run_k) - if __name__ == "__main__": Demo().run() diff --git a/nac3artiq/demo/min_artiq.py b/nac3artiq/demo/min_artiq.py index 06c3640d6..c7913b7a2 100644 --- a/nac3artiq/demo/min_artiq.py +++ b/nac3artiq/demo/min_artiq.py @@ -1,6 +1,7 @@ -from inspect import isclass, getmodule +from inspect import isclass, getfullargspec, getmodule from functools import wraps import sys +from types import SimpleNamespace from numpy import int32, int64 import nac3artiq @@ -37,13 +38,20 @@ def extern(function): return function -def kernel(class_or_function): +def kernel(function_or_method): """Decorates a function or method to be executed on the core device.""" - register_module_of(class_or_function) - @wraps(class_or_function) - def device_only(*args, **kwargs): - raise RuntimeError("Kernels must not be called directly, use core.run(kernel_function) instead") - return device_only + register_module_of(function_or_method) + argspec = getfullargspec(function_or_method) + if argspec.args and argspec.args[0] == "self": + @wraps(function_or_method) + def run_on_core(self, *args, **kwargs): + fake_method = SimpleNamespace(__self__=self, __name__=function_or_method.__name__) + self.core.run(fake_method, *args, **kwargs) + else: + @wraps(function_or_method) + def run_on_core(*args, **kwargs): + raise RuntimeError("Kernel functions need explicit core.run()") + return run_on_core def portable(function):