Repository: bluesky/bluesky
Pull request #1661 (2024-03-05 21:39:20)

Pass `md_validator` a deep copy that it can normalize
This PR adds a temporary fix to let us run the kafka consumers that run the prefect end of run workflow from the beamline



Patch:
diff --git a/bluesky/plans.py b/bluesky/plans.py
index ef7dbda8..f948d23a 100644
--- a/bluesky/plans.py
+++ b/bluesky/plans.py
@@ -290,13 +290,13 @@ def list_grid_scan(detectors, *args, snake_axes=False, per_step=None, md=None):
                           for motor, pos_list in partition(2, args)),
            'extents': tuple([min(pos_list), max(pos_list)]
                             for motor, pos_list in partition(2, args)),
-           'snake_axes': snake_axes,
+           'snake_axes': repr(snake_axes),
            'plan_args': {'detectors': list(map(repr, detectors)),
                          'args': md_args,
                          'per_step': repr(per_step)},
            'plan_name': 'list_grid_scan',
            'plan_pattern': 'outer_list_product',
-           'plan_pattern_args': dict(args=md_args, snake_axes=snake_axes),
+           'plan_pattern_args': dict(args=md_args, snake_axes=repr(snake_axes)),
            'plan_pattern_module': plan_patterns.__name__,
            'motors': tuple(motor_names),
            'hints': {},
diff --git a/bluesky/run_engine.py b/bluesky/run_engine.py
index 311b82ed..77aca7bf 100644
--- a/bluesky/run_engine.py
+++ b/bluesky/run_engine.py
@@ -18,6 +18,7 @@ from .protocols import (Flyable, Locatable, Movable, Pausable, Preparable, Reada
                         Status, Stoppable, Triggerable, check_supports)
 
 import concurrent
+import copy
 
 from event_model import DocumentNames
 from .log import logger, msg_logger, state_logger, ComposableLogAdapter
@@ -226,6 +227,16 @@ class RunEngine:
         completely up to the user. The function's return value is
         ignored.
 
+    md_normalizer : callable, optional
+        a function that, similar to md_validator, raises and prevents starting
+        a run if it deems the metadata to be invalid or incomplete.
+        If it succeeds, it returns the normalized/transformed version of
+        the original metadata.
+        Expected signature: f(md)
+        Function should raise if md is invalid. What that means is
+        completely up to the user.
+        Expected return: normalized metadata
+
     scan_id_source : callable, optional
         a function that will be used to calculate scan_id. Default is to
         increment scan_id by 1 each time. However you could pass in a
@@ -358,7 +369,7 @@ class RunEngine:
 
     def __init__(self, md=None, *, loop=None, preprocessors=None,
                  context_managers=None, md_validator=None,
-                 scan_id_source=default_scan_id_source,
+                 md_normalizer=None, scan_id_source=default_scan_id_source,
                  during_task=None, call_returns_result=False):
         if loop is None:
             loop = asyncio.new_event_loop()
@@ -414,6 +425,9 @@ class RunEngine:
         if md_validator is None:
             md_validator = _default_md_validator
         self.md_validator = md_validator
+        if md_normalizer is None:
+            md_normalizer = _default_md_normalizer
+        self.md_normalizer = md_normalizer
         self.scan_id_source = scan_id_source
 
         self.max_depth = None
@@ -1776,12 +1790,13 @@ class RunEngine:
                        'plan_name': plan_name},
                       self.md)  # stateful, persistent metadata
         # The metadata is final. Validate it now, at the last moment.
-        # Use copy for some reasonable (admittedly not total) protection
-        # against users mutating the md with their validator.
         self.md_validator(dict(md))
 
+        # Apply normalizer at the same level of the validator
+        validated = self.md_normalizer(copy.deepcopy(md))
+
         current_run = self._run_bundlers[run_key] = type(self).RunBundler(
-            md, self.record_interruptions, self.emit, self.emit_sync, self.log,
+            validated, self.record_interruptions, self.emit, self.emit_sync, self.log,
             strict_pre_declare=self._require_stream_declaration
         )
 
@@ -2707,6 +2722,10 @@ def _default_md_validator(md):
             "BAD: sample=[1, 2] ")
 
 
+def _default_md_normalizer(md):
+    return md
+
+
 def _ensure_event_loop_running(loop):
     """
     Run an asyncio event loop forever on a background thread.


Test patch:
diff --git a/bluesky/tests/test_metadata.py b/bluesky/tests/test_metadata.py
index 6bbbf17b..723efab1 100644
--- a/bluesky/tests/test_metadata.py
+++ b/bluesky/tests/test_metadata.py
@@ -1,4 +1,5 @@
 import bluesky
+from bluesky.plans import count
 import ophyd
 
 
@@ -8,3 +9,61 @@ def test_blueskyversion(RE):
 
 def test_ophydversion(RE):
     assert RE.md['versions'].get('ophyd') == ophyd.__version__
+
+
+def test_old_md_validator(RE):
+    """
+    Test an old-style md_validator.
+
+    In older versions of bluesky we ignored the result of md_validator.
+    All it could do was raise; it could not provide a normalized valid
+    result.
+
+    When an md_validator returns None, we should use the input dict.
+    """
+
+    def md_validator(md):
+        return None
+
+    start_doc = None
+
+    def test_callback(name, doc):
+        nonlocal start_doc
+        if name == "start":
+            start_doc = doc
+
+    RE.md_validator = md_validator
+    RE(count([], md={"test": 1}), test_callback)
+    assert start_doc["test"] == 1
+
+
+def test_md_mormalizer(RE):
+    """
+    Test the new md_normalizer.
+
+    In current versions of bluesky, md_normalizer is introduced to run
+    alongside md_validator. It returns a normalized valid result,
+    or raises if it is unable to provide one.
+
+    When an md_normalizer returns a dict, we should use that.
+    """
+
+    def md_normalizer(md):
+        "Ensure top-level keys are lowercase."
+        return {key.lower(): value for key, value in md.items()}
+
+    metadata = {"TEST": 1, "a": {"b": {"c": [1]}}}
+    start_doc = None
+
+    def test_callback(name, doc):
+        nonlocal start_doc
+        if name == "start":
+            start_doc = doc
+
+    RE.md_normalizer = md_normalizer
+    RE(count([], md=metadata), test_callback)
+    assert "TEST" in metadata
+    assert "TEST" not in start_doc
+    assert "test" in start_doc
+    assert start_doc["test"] == 1
+    assert start_doc["a"]["b"]["c"] is not metadata["a"]["b"]["c"]