Skip to content

Commit

Permalink
MNT #893 add template support
Browse files Browse the repository at this point in the history
  • Loading branch information
prjemian committed Dec 2, 2023
1 parent d957808 commit f81fe38
Showing 1 changed file with 86 additions and 0 deletions.
86 changes: 86 additions & 0 deletions apstools/callbacks/nexus_writer.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@
"""

import datetime
import json
import logging
import pathlib
import time
Expand Down Expand Up @@ -94,6 +95,7 @@ def my_plan(dets, n=5):
~create_NX_group
~get_sample_title
~get_stream_link
~resolve_class_path
~wait_writer
~wait_writer_plan_stub
~write_data
Expand All @@ -108,6 +110,7 @@ def my_plan(dets, n=5):
~write_slits
~write_source
~write_streams
~write_templates
~write_user
New with apstools release 1.3.0.
Expand All @@ -123,6 +126,9 @@ def my_plan(dets, n=5):
nxdata_signal_axes = None # name of dataset for X axis on plot
root = None # instance of h5py.File

template_key = "nxwriter_template" # TODO: refactor template from dict to list
"""The template (dict) is written as a JSON string to this metadata key."""

_external_file_read_timeout = 20
_external_file_read_retry_delay = 0.5
_writer_active = False
Expand Down Expand Up @@ -246,6 +252,34 @@ def h5string(self, text):
text = text or ""
return text.encode("utf8")

def resolve_class_path(self, class_path):
"""
Parse the class path, make any groups, return the HDF5 address.
New with apstools release 1.6.18.
"""
addr = ""
for level in class_path.split("/"):
if ":" in level:
group_name, nx_class = level.split(":")
if not nx_class.startswith("NX"):
raise ValueError(f"nx_class must start with 'NX'. Received {nx_class=!r}")
if group_name not in self.root[addr]:
# fmt: off
logger.info(
"make HDF5 group with @NX_class=%r at address '%s/%s'",
nx_class, addr, group_name
)
# fmt: on
self.create_NX_group(self.root[addr], level)
addr += f"/{group_name}"
else:
addr += f"/{level}"
addr = addr.replace("//", "/")

logger.debug("HDF5 address=%r", addr)
return addr

def wait_writer(self):
"""
Wait for the writer to finish. For interactive use (Not in a plan).
Expand Down Expand Up @@ -414,6 +448,11 @@ def write_entry(self):
nxentry["plan_name"] = self.root["/entry/instrument/bluesky/metadata/plan_name"]
nxentry["entry_identifier"] = self.root["/entry/instrument/bluesky/uid"]

try:
self.write_templates()
except Exception as exc:
logger.warning("Problem writing template(s): %s", exc)

return nxentry

def write_instrument(self, parent):
Expand Down Expand Up @@ -753,6 +792,53 @@ def write_streams(self, parent):

return bluesky

def write_templates(self):
"""
Process any link templates provided as run metadata.
New in v1.6.18.
"""
addr = f"/entry/instrument/bluesky/metadata/{self.template_key}"
if addr not in self.root:
return
templates = json.loads(self.root[addr][()])

for source, target in templates:
if "/@" in source:
p = source.rfind("/")
if source.find("/@") != p: # more than one match
raise ValueError(f"Only one attribute can be named. Received: {source!r}")
h5addr = self.resolve_class_path(source[:p])
attr = source[p + 2 :]
logger.debug("Set attribute: group=%r attr=%r value=%r", h5addr, attr, target)
if h5addr in self.root:
self.root[h5addr].attrs[attr] = target
else:
logger.warning("group %r not in root %r", h5addr, self.root.name)
elif source.endswith("="):
p = source.rfind("/")
h5addr = self.resolve_class_path(source[:p])
field = source.split("/")[-1].rstrip("=")
if h5addr in self.root:
if isinstance(target, (int, float)):
target = [target]
ds = self.root[h5addr].create_dataset(field, data=target)
ds.attrs["target"] = ds.name
# fmt: off
logger.info(
"Set constant field: group=%r field=%r value=%r",
h5addr, field, ds[()]
)
# fmt: on
else:
logger.warning("group %r not in root %r", h5addr, self.root.name)
elif source in self.root:
h5addr = self.resolve_class_path(target)
self.root[h5addr] = self.root[source]
logger.debug("Template: Linked %r to %r", source, h5addr)
else:
logger.warning("Not handled: source=%r target=%r", source, target)

def write_user(self, parent):
"""
group: /entry/contact:NXuser
Expand Down

0 comments on commit f81fe38

Please sign in to comment.