From 868ead38c76e1c049674cf3ea6ffacb31fbca0c2 Mon Sep 17 00:00:00 2001
From: Jan Gosmann <jgosmann@uwaterloo.ca>
Date: Mon, 30 Apr 2018 16:07:05 -0400
Subject: [PATCH] Implement hooks with register functions/decorators

Addresses #953.
---
 nengo_gui/components/sim_control.py |  5 ++++
 nengo_gui/guibackend.py             |  1 +
 nengo_gui/hooks.py                  | 40 +++++++++++++++++++++++++++++
 nengo_gui/page.py                   |  9 ++++++-
 4 files changed, 54 insertions(+), 1 deletion(-)
 create mode 100644 nengo_gui/hooks.py

diff --git a/nengo_gui/components/sim_control.py b/nengo_gui/components/sim_control.py
index b05a80c1..1f2abe5a 100644
--- a/nengo_gui/components/sim_control.py
+++ b/nengo_gui/components/sim_control.py
@@ -8,6 +8,7 @@
 
 from nengo_gui.components.component import Component
 import nengo_gui.exec_env
+from nengo_gui import hooks
 from nengo_gui.server import WebSocketFrame
 
 
@@ -190,11 +191,15 @@ def javascript(self):
     def message(self, msg):
         if msg == 'pause':
             self.paused = True
+            hooks.on_pause.execute(self.page)
         elif msg == 'config':
             self.send_config_options = True
         elif msg == 'continue':
             if self.page.sim is None:
                 self.page.rebuild = True
+                hooks.on_start.execute(self.page)
+            else:
+                hooks.on_continue.execute(self.page)
             self.paused = False
         elif msg == 'reset':
             self.paused = True
diff --git a/nengo_gui/guibackend.py b/nengo_gui/guibackend.py
index 57b1ded1..ffccda4c 100644
--- a/nengo_gui/guibackend.py
+++ b/nengo_gui/guibackend.py
@@ -432,6 +432,7 @@ def create_page(self, filename, reset_cfg=False):
 
     def remove_page(self, page):
         self._last_access = time.time()
+        page.close()
         self.pages.remove(page)
         if (not self._shutting_down and self.settings.auto_shutdown > 0 and
                 len(self.pages) <= 0):
diff --git a/nengo_gui/hooks.py b/nengo_gui/hooks.py
new file mode 100644
index 00000000..39514a09
--- /dev/null
+++ b/nengo_gui/hooks.py
@@ -0,0 +1,40 @@
+from collections import defaultdict
+
+
+class Context(object):
+    context_stack = [None]
+
+    def __init__(self, page):
+        self.page = page
+
+    def __enter__(self):
+        self.context_stack.append(self.page)
+        for hook in (on_step, on_start, on_pause, on_continue, on_close):
+            # NOTE entering context clears hooks, this might be surprising and
+            # there might be a better solution to prevent accumulation of hooks
+            # from multiple code executions.
+            if self.page in hook.callbacks:
+                del hook.callbacks[self.page]
+        return self
+
+    def __exit__(self, exc_type, exc_value, tb):
+        assert self.context_stack.pop() is self.page
+
+
+class Hook(object):
+    def __init__(self):
+        self.callbacks = defaultdict(list)  # FIXME page references should be weak
+
+    def __call__(self, fn):
+        self.callbacks[Context.context_stack[-1]].append(fn)
+
+    def execute(self, page):
+        for cb in self.callbacks[page] + self.callbacks[None]:
+            cb(page.sim)
+
+
+on_step = Hook()
+on_start = Hook()
+on_pause = Hook()
+on_continue = Hook()
+on_close = Hook()
diff --git a/nengo_gui/page.py b/nengo_gui/page.py
index a5a99801..5ec7d3a2 100644
--- a/nengo_gui/page.py
+++ b/nengo_gui/page.py
@@ -10,6 +10,7 @@
 import nengo
 
 import nengo_gui
+from nengo_gui import hooks
 import nengo_gui.user_action
 import nengo_gui.config
 import nengo_gui.seed_generation
@@ -214,7 +215,7 @@ def execute(self, code):
 
         exec_env = nengo_gui.exec_env.ExecutionEnvironment(self.filename)
         try:
-            with exec_env:
+            with exec_env, hooks.Context(self):
                 compiled = compile(code, nengo_gui.exec_env.compiled_filename,
                                    'exec')
                 exec(compiled, code_locals)
@@ -501,6 +502,7 @@ def runner(self):
                         self.sim.run_steps(self.sim.max_steps)
                     else:
                         self.sim.step()
+                        hooks.on_step.execute(self)
                 except Exception as err:
                     if self.finished:
                         return
@@ -508,8 +510,13 @@ def runner(self):
                     self.error = dict(trace=traceback.format_exc(), line=line)
                     self.sim = None
             while self.sims_to_close:
+                hooks.on_close.execute(self)
                 self.sims_to_close.pop().close()
 
             if self.rebuild:
                 self.build()
         self.sim = None
+
+    def close(self):
+        if self.sim is not None:
+            hooks.on_close.execute(self)