plan_mutator_test.py

from itertools import tee
import bluesky.plans as bp
import bluesky.plan_stubs as bps
import bluesky.preprocessors as bpp
from bluesky.run_engine import RunEngine
from bluesky.simulators import summarize_plan

from ophyd.sim import hw

ns = hw()
RE = RunEngine()
RE.msg_hook = print

def _close_shutter():
    yield from bps.abs_set(ns.motor, 0, wait=True)


def _open_shutter():
    yield from bps.abs_set(ns.motor, 1, wait=True)


def shutter_control_full(msg):
    if msg.command=='trigger':
        def inner():
            yield from _open_shutter()
            yield from bpp.single_gen(msg)
        return inner(), None
    elif msg.command=='save':
        return None, _close_shutter()
    else:
        return None, None

    
def shutter_control_full_broken(msg):
    if msg.command=='trigger':
        def inner():
            yield from _open_shutter()
            yield from msg
        return inner(), None
    elif msg.command=='save':
        return None, _close_shutter()
    else:
        return None, None
    

def shutter_control_half(msg):
    if msg.command=='save':
        return None, _close_shutter()
    else:
        return None, None

p = bp.count([ns.det, ns.motor], 5, 0.5)
q = bp.count([ns.det, ns.motor], 5, 0.5)


full_p = bpp.plan_mutator(p, shutter_control_full)
borken_p = bpp.plan_mutator(q, shutter_control_full_broken)

# this works fine
print(list(full_p))

# this raises error : str doesn't have attribute "command"
# comes from the fact "trigger" is not a Msg object anymore.
print(list(broken_p))