diff --git a/src/scippnexus/base.py b/src/scippnexus/base.py index a6c16ca0..53b3a5f1 100644 --- a/src/scippnexus/base.py +++ b/src/scippnexus/base.py @@ -9,7 +9,7 @@ from functools import lru_cache from pathlib import PurePosixPath from types import MappingProxyType -from typing import Any, overload +from typing import Any, TypeGuard, overload import h5py as h5 import numpy as np @@ -153,6 +153,15 @@ class NXroot(NXobject): pass +def _is_nx_class_selector(sel: Any) -> TypeGuard[type | list[type]]: + def isclass(x: Any) -> bool: + return inspect.isclass(x) and issubclass(x, Field | NXobject) + + return isclass(sel) or ( + isinstance(sel, list) and len(sel) > 0 and all(isclass(x) for x in sel) + ) + + class Group(Mapping): """ A group in a NeXus file. @@ -304,16 +313,36 @@ def items(self) -> ItemsView[str, Field | Group]: """Return a view of pairs of the keys and the group's elements.""" return self._children.items() + def __contains__(self, key: object) -> bool: + """Check for a child name/path or at least one matching NeXus class.""" + if _is_nx_class_selector(key): + return next(self._names_by_nx_class(key), None) is not None + return super().__contains__(key) + + def get(self, key: str | type | list[type], default: Any = None) -> Any: + """Return a child or class selection, or default when there is no match. + + Unlike ``group[NXlog]``, which returns an empty dict if no child matches, + ``group.get(NXlog)`` returns ``default`` in that case. This keeps ``get`` + consistent with ``NXlog in group``. + """ + if _is_nx_class_selector(key): + return self._get_children_by_nx_class(key) or default + return super().get(key, default) + + def _names_by_nx_class(self, select: type | list[type]) -> Iterator[str]: + # Does not populate fields, so checking membership stays cheap. + selectors = (select,) if isinstance(select, type) else select + return ( + name + for name, child in self._children.items() + if (Field if isinstance(child, Field) else child.nx_class) in selectors + ) + def _get_children_by_nx_class( self, select: type | list[type] ) -> dict[str, NXobject | Field]: - children = {} - select = (select,) if isinstance(select, type) else select - for key, child in self._children.items(): - nx_class = Field if isinstance(child, Field) else child.nx_class - if nx_class is not None and any(nx_class == sel for sel in select): - children[key] = self[key] - return children + return {name: self[name] for name in self._names_by_nx_class(select)} @overload def __getitem__(self, sel: str) -> Group | Field: ... @@ -335,9 +364,10 @@ def __getitem__(self, sel): - String name: The child group or child dataset of that name is returned. - Class such as ``NXdata`` or ``NXlog``: A dict containing all direct children - with a matching ``NX_class`` attribute are returned. Also accepts a tuple of + with a matching ``NX_class`` attribute are returned. Also accepts a list of classes. ``Field`` selects all child fields, i.e., all datasets but not - groups. + groups. Returns an empty dict if no children match. Use ``NXlog in group`` + or ``group.get(NXlog)`` to check whether any child matches. - Scipp-style index: Load the specified slice of the current group, returning a :class:`scipp.DataArray` or :class:`scipp.DataGroup`. @@ -372,12 +402,7 @@ def __getitem__(self, sel): self._populate_fields() return child - def isclass(x): - return inspect.isclass(x) and issubclass(x, Field | NXobject) - - if isclass(sel) or ( - isinstance(sel, list) and len(sel) and all(isclass(x) for x in sel) - ): + if _is_nx_class_selector(sel): return self._get_children_by_nx_class(sel) dg = self._nexus.read_children(sel) diff --git a/tests/nexus_test.py b/tests/nexus_test.py index 110dabec..884fe79b 100644 --- a/tests/nexus_test.py +++ b/tests/nexus_test.py @@ -133,18 +133,26 @@ def test_nxobject_getitem_by_class(nxroot) -> None: nxroot['entry'].create_class('events_1', NXevent_data) assert list(nxroot[NXentry]) == ['entry'] assert list(nxroot[NXmonitor]) == ['monitor'] - assert list(nxroot['entry'][NXmonitor]) == [] # not nested - assert list(nxroot[NXlog]) == [] # nested + assert nxroot['entry'][NXmonitor] == {} # not nested + assert nxroot[NXlog] == {} # nested assert list(nxroot['entry'][NXlog]) == ['log'] assert set(nxroot['entry'][NXevent_data]) == {'events_0', 'events_1'} +def test_nxobject_contains_by_class(nxroot) -> None: + nxroot.create_class('monitor', NXmonitor) + + assert NXentry in nxroot + assert NXmonitor in nxroot + assert NXlog not in nxroot + + def test_nxobject_getitem_by_class_get_fields(nxroot) -> None: nxroot['entry'].create_class('log', NXlog) nxroot['entry'].create_class('events_0', NXevent_data) nxroot['entry']['field1'] = sc.arange('event', 4.0, unit='ns') nxroot['entry']['field2'] = sc.arange('event', 2.0, unit='ns') - assert list(nxroot[snx.Field]) == [] + assert nxroot[snx.Field] == {} assert set(nxroot['entry'][snx.Field]) == {'field1', 'field2'} @@ -161,6 +169,73 @@ def test_nxobject_getitem_by_class_list(nxroot) -> None: assert set(nxroot['entry'][[NXlog, snx.Field]]) == {'log', 'field1'} +@pytest.mark.parametrize('select', [NXlog, snx.Field, [NXlog, snx.Field]]) +def test_nxobject_empty_class_selection_preserves_getitem(nxroot, select) -> None: + default = object() + + assert nxroot[select] == {} + assert select not in nxroot + assert nxroot.get(select) is None + assert nxroot.get(select, default) is default + + +@pytest.mark.parametrize('select', [NXentry, [NXentry, NXlog]]) +def test_nxobject_present_class_selection_supports_contains_and_get( + nxroot, select +) -> None: + assert select in nxroot + assert set(nxroot.get(select)) == {'entry'} + assert nxroot.get(select)['entry'] is nxroot['entry'] + + +def test_nxobject_class_selection_matches_exact_class_not_base_class(nxroot) -> None: + assert nxroot[snx.NXobject] == {} + assert snx.NXobject not in nxroot + assert nxroot.get(snx.NXobject) is None + + +def test_nxobject_class_membership_does_not_populate_fields( + nxroot, monkeypatch +) -> None: + nxroot['entry']['field'] = sc.arange('x', 3.0) + entry = nxroot['entry'] + + def fail_populate_fields(self): + pytest.fail('Class membership must not populate fields') + + monkeypatch.setattr(snx.Group, '_populate_fields', fail_populate_fields) + + assert snx.Field in entry + assert [NXlog, snx.Field] in entry + assert NXlog not in entry + + +def test_nxobject_contains_and_get_preserve_string_paths_and_empty_groups( + nxroot, +) -> None: + nxroot['entry'].create_class('log', NXlog) + default = object() + + for path in ('entry', 'entry/log', '/entry/log'): + assert path in nxroot + assert nxroot.get(path).name == '/' + path.lstrip('/') + for path in ('absent', 'entry/absent', '/entry/absent'): + assert path not in nxroot + assert nxroot.get(path, default) is default + + +def test_nxobject_get_returns_empty_group_instead_of_default( + nxroot, +) -> None: + nxroot.create_class('log', NXlog) + log = nxroot['log'] + default = object() + + # An empty Group is falsey, but a successful get must not use the default. + assert len(log) == 0 + assert nxroot.get('log', default) is log + + def test_nxobject_dataset_items_are_returned_as_Field(nxroot) -> None: events = nxroot['entry'].create_class('events_0', NXevent_data) events['event_time_offset'] = sc.arange('event', 5)