feat/auto-refresh stale at_each_angle hooks on redefinition

register_at_each_angle_hook() stores a function object by reference, so
editing a hook's def and forgetting to re-register left the old version
running silently -- easy to miss mid-beamtime. Detect it via
func.__globals__ (a redefined top-level def, or a reloaded module,
mutates the same namespace dict in place) and auto-adopt the new
definition with a printed note instead of failing or staying stale.

Centralizes the shared lookup/error logic (previously duplicated in
lamni.py and flomni.py) into TomoQueueMixin._resolve_at_each_angle_hook().

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

csaxs_bec/bec_ipython_client/plugins/LamNI/lamni.py

@@ -1174,14 +1174,7 @@ class LamNI(TomoQueueMixin, LamNIAlignmentMixin, LamNIOpticsMixin, LamniGuiTools
     def _at_each_angle(self, angle: float) -> None:
         hook_name = self.at_each_angle_hook
         if hook_name:
-            hook = self._at_each_angle_hooks.get(hook_name)
-            if hook is None:
-                raise LamNIError(
-                    f"at_each_angle_hook '{hook_name}' is not registered in this "
-                    "session. Hooks are session-only and do not survive a kernel "
-                    f"restart -- call register_at_each_angle_hook({hook_name!r}, "
-                    "<func>) again, then re-run tomo_queue_execute() to resume."
-                )
+            hook = self._resolve_at_each_angle_hook(hook_name, LamNIError)
             hook(self, angle)
             return
 

csaxs_bec/bec_ipython_client/plugins/OMNY_shared/tomo_queue_mixin.py

@@ -40,6 +40,7 @@ import datetime
 import inspect
 import json
 import uuid
+from typing import Callable
 
 from typeguard import check_type
 
@@ -248,6 +249,15 @@ class TomoQueueMixin:
         time if the name isn't registered in the session running the queue --
         re-run this registration call, then resume ``tomo_queue_execute()``.
 
+        If you edit ``func``'s source and re-run the ``def`` (in a cell, or
+        via ``importlib.reload()`` of the module it lives in) without calling
+        this again, the next execution automatically picks up the new
+        definition and prints a note saying so -- see
+        ``_resolve_at_each_angle_hook()``. This does NOT cover plain
+        ``%run script.py`` (without ``-i``), which re-executes in a fresh
+        namespace each time; use ``%run -i`` or ``import`` + ``reload()`` for
+        hooks defined in a file.
+
         Example -- record a projection, insert a polarizer, record again::
 
             def polarizer_modulation(setup, angle):
@@ -267,6 +277,40 @@ class TomoQueueMixin:
         self._publish_at_each_angle_hooks()
         print(f"Registered at_each_angle hook '{name}'.")
 
+    def _resolve_at_each_angle_hook(self, hook_name: str, error_cls: type) -> Callable:
+        """Look up a registered at_each_angle hook by name for execution,
+        raising ``error_cls`` if it isn't registered in this session.
+
+        If the function has been redefined since it was registered -- i.e.
+        the name it was defined under, in the namespace it was defined in
+        (``func.__globals__``), no longer points at the same object -- the
+        live definition is adopted in its place and a note is printed. This
+        catches editing and re-running a ``def`` in the same session/cell, or
+        in a file that's then ``importlib.reload()``-ed (the module dict is
+        ``func.__globals__`` itself, mutated in place by reload); it does not
+        catch plain ``%run`` (without ``-i``), which builds a fresh namespace
+        per run.
+        """
+        hook = self._at_each_angle_hooks.get(hook_name)
+        if hook is None:
+            raise error_cls(
+                f"at_each_angle_hook '{hook_name}' is not registered in this "
+                "session. Hooks are session-only and do not survive a kernel "
+                f"restart -- call register_at_each_angle_hook({hook_name!r}, "
+                "<func>) again, then re-run tomo_queue_execute() to resume."
+            )
+        live = getattr(hook, "__globals__", {}).get(getattr(hook, "__name__", None))
+        if callable(live) and live is not hook:
+            self._at_each_angle_hooks[hook_name] = live
+            print(
+                f"Note: at_each_angle_hook '{hook_name}' was redefined since "
+                "it was registered -- using the new definition. Call "
+                "register_at_each_angle_hook() explicitly if that's not what "
+                "you want."
+            )
+            hook = live
+        return hook
+
     def unregister_at_each_angle_hook(self, name: str) -> None:
         """Remove a previously registered at_each_angle hook by name."""
         if self._at_each_angle_hooks.pop(name, None) is None:

csaxs_bec/bec_ipython_client/plugins/flomni/flomni.py

@@ -2853,14 +2853,7 @@ class Flomni(
     def _at_each_angle(self, angle: float) -> None:
         hook_name = self.at_each_angle_hook
         if hook_name:
-            hook = self._at_each_angle_hooks.get(hook_name)
-            if hook is None:
-                raise FlomniError(
-                    f"at_each_angle_hook '{hook_name}' is not registered in this "
-                    "session. Hooks are session-only and do not survive a kernel "
-                    f"restart -- call register_at_each_angle_hook({hook_name!r}, "
-                    "<func>) again, then re-run tomo_queue_execute() to resume."
-                )
+            hook = self._resolve_at_each_angle_hook(hook_name, FlomniError)
             hook(self, angle)
             return
 

docs/user/ptychography/flomni.md

@@ -515,9 +515,22 @@ again to resume.
 
 **Editing a hook after registering it:** `register_at_each_angle_hook()` stores
 whichever function object you pass it at that moment — it is not a live link to the
-function's name. If you edit the function's source and re-run the `def` in your
-session but do **not** call `register_at_each_angle_hook()` again, the *old* version
-keeps running. Always re-register after an edit.
+function's name. If you edit the function's source and re-run the `def` (in a cell,
+or reload the file it's defined in — see below), the change is picked up
+automatically the next time the hook runs, and a note is printed telling you so.
+You don't need to call `register_at_each_angle_hook()` again after an edit.
+
+**Recommended: define hooks in a file, not inline in the shell.** For anything
+beyond a quick one-off, put the hook function in a `.py` file — ideally together
+with the `tomo_queue_add()`/`at_each_angle_hook` calls that use it, so the file is a
+single, compact record of exactly what ran during the experiment. Load it with
+`import myhooks` (or `from myhooks import my_hook`) and, after editing, reload with
+`importlib.reload(myhooks)` — the automatic pick-up described above still works,
+because it detects the change via the module's own namespace, not the local name
+you imported it under. `%run -i myhooks.py` also works (runs in the current
+namespace, same as a cell). **Avoid plain `%run myhooks.py`** (without `-i`): it
+executes in a fresh, throwaway namespace each time, so an edit-and-rerun is *not*
+picked up automatically and the old version keeps running silently.
 
 **`tomo_scan_projection()` vs `tomo_acquire_at_angle()`:** these are not
 interchangeable. `tomo_scan_projection(angle)` always runs a full Fermat-scan

docs/user/ptychography/lamni.md

@@ -282,9 +282,22 @@ again to resume.
 
 **Editing a hook after registering it:** `register_at_each_angle_hook()` stores
 whichever function object you pass it at that moment — it is not a live link to the
-function's name. If you edit the function's source and re-run the `def` in your
-session but do **not** call `register_at_each_angle_hook()` again, the *old* version
-keeps running. Always re-register after an edit.
+function's name. If you edit the function's source and re-run the `def` (in a cell,
+or reload the file it's defined in — see below), the change is picked up
+automatically the next time the hook runs, and a note is printed telling you so.
+You don't need to call `register_at_each_angle_hook()` again after an edit.
+
+**Recommended: define hooks in a file, not inline in the shell.** For anything
+beyond a quick one-off, put the hook function in a `.py` file — ideally together
+with the `tomo_queue_add()`/`at_each_angle_hook` calls that use it, so the file is a
+single, compact record of exactly what ran during the experiment. Load it with
+`import myhooks` (or `from myhooks import my_hook`) and, after editing, reload with
+`importlib.reload(myhooks)` — the automatic pick-up described above still works,
+because it detects the change via the module's own namespace, not the local name
+you imported it under. `%run -i myhooks.py` also works (runs in the current
+namespace, same as a cell). **Avoid plain `%run myhooks.py`** (without `-i`): it
+executes in a fresh, throwaway namespace each time, so an edit-and-rerun is *not*
+picked up automatically and the old version keeps running silently.
 
 **GUI:** the tomo parameters panel has an "At-each-angle hook" dropdown, listing
 whatever is currently registered from the CLI (`register_at_each_angle_hook()`) —

tests/tests_bec_ipython_client/test_lamni_tomo_queue.py

@@ -204,3 +204,75 @@ def test_at_each_angle_raises_for_unregistered_hook_name():
     except LamNIError as exc:
         assert "missing_hook" in str(exc)
         assert "not registered" in str(exc)
+
+
+def test_at_each_angle_hook_refreshed_when_redefined(capsys):
+    """Redefining a hook's `def` in the same namespace (as happens re-running
+    a cell in IPython) and forgetting to re-register must not silently keep
+    running the old version -- the live definition should be picked up
+    automatically, with a note printed."""
+    lamni = make_lamni()
+    ns = {}
+    exec("def hook_func(setup, angle):\n    setup.marker = 'old'\n", ns)
+    lamni.register_at_each_angle_hook("swap", ns["hook_func"])
+    lamni.at_each_angle_hook = "swap"
+
+    exec("def hook_func(setup, angle):\n    setup.marker = 'new'\n", ns)
+    lamni._at_each_angle(5.0)
+
+    assert lamni.marker == "new"
+    assert "redefined" in capsys.readouterr().out
+
+
+def test_register_at_each_angle_hook_with_lambda_never_flagged_stale(capsys):
+    """Lambdas have no real namespace binding under their `__name__`
+    ("<lambda>"), so the redefinition check must never fire a false positive
+    for them -- covered separately from the dispatch test above since that
+    one doesn't assert anything about stdout."""
+    lamni = make_lamni()
+    lamni.register_at_each_angle_hook("record", lambda self, angle: None)
+    lamni.at_each_angle_hook = "record"
+
+    lamni._at_each_angle(1.0)
+
+    assert "redefined" not in capsys.readouterr().out
+
+
+def test_at_each_angle_hook_refreshed_after_module_reload(tmp_path, capsys):
+    """Mirrors the recommended file-based workflow: define the hook in a
+    .py file, load it, edit the file, and reload the module -- the module's
+    __dict__ (== the hook function's __globals__) is mutated in place by
+    importlib.reload(), so this must be picked up the same way an
+    interactive redefinition is."""
+    import importlib
+    import os
+    import sys
+
+    module_name = "lamni_hook_reload_test_module"
+    module_path = tmp_path / f"{module_name}.py"
+
+    def write(marker_value, mtime):
+        module_path.write_text(f"def hook_func(setup, angle):\n    setup.marker = {marker_value!r}\n")
+        os.utime(module_path, (mtime, mtime))
+
+    # Distinct, well-separated mtimes -- otherwise both writes can land in
+    # the same filesystem mtime-resolution window and the loader reuses the
+    # cached .pyc from the first write, silently defeating the reload.
+    write("old", 1_700_000_000)
+    sys.path.insert(0, str(tmp_path))
+    try:
+        module = importlib.import_module(module_name)
+        lamni = make_lamni()
+        lamni.register_at_each_angle_hook("swap", module.hook_func)
+        lamni.at_each_angle_hook = "swap"
+
+        write("new", 1_700_000_010)
+        importlib.reload(module)
+
+        lamni._at_each_angle(5.0)
+
+        assert lamni.marker == "new"
+        assert "redefined" in capsys.readouterr().out
+    finally:
+        sys.path.remove(str(tmp_path))
+        sys.modules.pop(module_name, None)