identify if data is missing in util.restore_energy()

hkl/tests/test_save_restore_UB.py

@@ -5,6 +5,7 @@
 from hkl import SimulatedK4CV
 from hkl.calc import A_KEV
 from ophyd.sim import hw
+import pandas as pd
 import bluesky.plans as bp
 import databroker
 import hkl.util
@@ -157,7 +158,7 @@ def scans():
     assert runs.diffractometer_name.to_list() == "fourc kappa fourc kappa".split()
 
 
-def test_no_primary_stream(cat, RE, fourc, kappa):
+def test_no_primary_stream(cat, RE, fourc):
     det = hw().noisy_det
 
     def my_plan():
@@ -171,12 +172,42 @@ def scans():
         yield from bp.count([fourc])
         yield from my_plan()
 
-    RE(scans())
+    uids = RE(scans())
+    assert len(uids) == 2
+
     runs = hkl.util.list_orientation_runs(cat)
     # my_plan() has no primary stream
     assert len(runs.scan_id) == 1
 
 
+def test_missing_energy_key(cat, RE, fourc):
+    """Issue 216."""
+
+    def scans():
+        yield from bp.count([fourc])
+
+    uids = RE(scans())
+    assert len(uids) == 1
+    assert uids[0] in cat
+
+    runs = hkl.util.list_orientation_runs(cat)
+    assert isinstance(runs, pd.DataFrame)
+
+    run = cat[uids[0]]
+    orientations = hkl.util.run_orientation_info(run)
+    assert len(orientations) == 1
+
+    assert fourc.name in orientations
+    orientation = orientations[fourc.name]
+    assert "energy" in orientation
+
+    # trigger the error by removing the "energy" key
+    with pytest.raises(KeyError) as exinfo:
+        orientation.pop("energy")
+        hkl.util.restore_energy(orientation, fourc)
+    assert " Cannot restore diffractometer energy " in str(exinfo.value)
+
+
 def test_restore_orientation(cat, RE, fourc):
     RE(bp.count([fourc]))
     fourc_orient = hkl.util.run_orientation_info(cat[-1])["fourc"]

hkl/util.py

@@ -385,8 +385,22 @@ def restore_energy(orientation, diffractometer):
     diffractometer : :class:`~hkl.diffract.Diffractometer()`
         Diffractometer object.
     """
-    for attr in "energy energy_units energy_offset".split():
-        _smart_signal_update(orientation[attr], getattr(diffractometer, attr))
+    # get _all_ the expected keys
+    try:
+        kv_dict = {
+            orientation[attr]: getattr(diffractometer, attr)
+            for attr in "energy energy_units energy_offset".split()
+        }
+    except KeyError as exc:
+        # fmt: off
+        raise KeyError(
+            f"{diffractometer.name}: Cannot restore "
+            f"diffractometer energy due to missing {exc} term."
+        )
+        # fmt: on
+    # update the signals
+    for k, v in kv_dict.items():
+        _smart_signal_update(k, v)
 
 
 def restore_reflections(orientation, diffractometer):