Add open_component_group

Author: jl-wynen

src/ess/reduce/nexus/__init__.py

@@ -20,6 +20,7 @@
     load_all_components,
     load_component,
     load_data,
+    open_component_group,
     open_nexus_file,
 )
 from .workflow import GenericNeXusWorkflow
@@ -32,6 +33,7 @@
     'load_all_components',
     'load_component',
     'load_data',
+    'open_component_group',
     'open_nexus_file',
     'types',
 ]

src/ess/reduce/nexus/_nexus_loader.py

@@ -46,29 +46,52 @@ def load_component(
     location: NeXusLocationSpec,
     *,
     nx_class: type[snx.NXobject],
+    parent_class: type[snx.NXobject] | None = None,
     definitions: Mapping | None | NoNewDefinitionsType = NoNewDefinitions,
 ) -> sc.DataGroup:
-    """Load a single component of a given class from NeXus."""
+    """Load a single component of a given class from NeXus.
+
+    Parameters
+    ----------
+    location:
+        Specifies (part of) the location of the component to load.
+    nx_class:
+        NeXus class of the component to load.
+    parent_class:
+        NeXus class of the parent of the component to load.
+        If ``None``, is deduced from ``nx_class`` if possible.
+    definitions:
+        Application definitions to use for the file.
+
+    Returns
+    -------
+    :
+        The loaded component as a data group.
+    """
     selection = location.selection
-    group_name = location.component_name
-    with _open_component_parent(
-        location, nx_class=nx_class, definitions=definitions
-    ) as parent:
-        component = _unique_child_group(parent, nx_class, group_name)
-        loaded = cast(sc.DataGroup, component[selection])
-        loaded['nexus_component_name'] = component.name.rsplit('/', 1)[-1]
+    with open_component_group(
+        location,
+        nx_class=nx_class,
+        parent_class=parent_class,
+        definitions=definitions,
+    ) as group:
+        loaded = cast(sc.DataGroup, group[selection])
+        loaded['nexus_component_name'] = group.name.rsplit('/', 1)[-1]
     return loaded
 
 
 def load_all_components(
     location: NeXusAllLocationSpec,
     *,
     nx_class: type[snx.NXobject],
+    parent_class: type[snx.NXobject] | None = None,
     definitions: Mapping | None | NoNewDefinitionsType = NoNewDefinitions,
 ) -> sc.DataGroup:
     """Load all components of a given class from NeXus."""
     with _open_component_parent(
-        location, nx_class=nx_class, definitions=definitions
+        location,
+        parent_class=_deduce_component_parent_class(nx_class, parent_class),
+        definitions=definitions,
     ) as parent:
         components = sc.DataGroup()
         for name, component in parent[nx_class].items():
@@ -226,21 +249,77 @@ def _open_nexus_file_from_path(
 
 
 @contextmanager
-def _open_component_parent(
+def open_component_group(
     location: NeXusLocationSpec,
     *,
     nx_class: type[snx.NXobject],
+    parent_class: type[snx.NXobject] | None = None,
+    definitions: Mapping | None | NoNewDefinitionsType = NoNewDefinitions,
+) -> Generator[snx.Group, None, None]:
+    """Open the HDF5 group of a NeXus component.
+
+    Parameters
+    ----------
+    location:
+        Specifies (part of) the location of the component to load.
+    nx_class:
+        NeXus class of the component to load.
+    parent_class:
+        NeXus class of the parent of the component to load.
+        If ``None``, is deduced from ``nx_class`` if possible.
+    definitions:
+        Application definitions to use for the file.
+
+    Returns
+    -------
+    :
+        A context manager for the group of the specified component.
+    """
+    group_name = location.component_name
+    with _open_component_parent(
+        location,
+        parent_class=_deduce_component_parent_class(nx_class, parent_class),
+        definitions=definitions,
+    ) as parent:
+        yield _unique_child_group(parent, nx_class, group_name)
+
+
+@contextmanager
+def _open_component_parent(
+    location: NeXusLocationSpec,
+    *,
+    parent_class: type[snx.NXobject],
     definitions: Mapping | None | NoNewDefinitionsType = NoNewDefinitions,
 ) -> Generator[snx.Group, None, None]:
     """Locate the parent group of a NeXus component."""
     file_path = location.filename
     entry_name = location.entry_name
     with open_nexus_file(file_path, definitions=definitions) as f:
         entry = _unique_child_group(f, snx.NXentry, entry_name)
-        if nx_class is snx.NXsample:
-            yield entry
-        else:
-            yield _unique_child_group(entry, snx.NXinstrument, None)
+        match parent_class:
+            case snx.NXentry:
+                yield entry
+            case snx.NXinstrument:
+                yield _unique_child_group(entry, snx.NXinstrument, None)
+            case _:
+                raise NotImplementedError(
+                    f"No support for loading a NeXus component from a {parent_class}."
+                )
+
+
+def _deduce_component_parent_class(
+    nx_class: type[snx.NXobject], parent_class: type[snx.NXobject] | None
+) -> type[snx.NXobject]:
+    if parent_class is not None:
+        return parent_class
+
+    match nx_class:
+        case snx.NXsample:
+            return snx.NXentry
+        case _:
+            # Most components are in the instrument,
+            # callers need to override this for specialized components stored elsewhere.
+            return snx.NXinstrument
 
 
 def _unique_child_group(