Fix NX class membership checks

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)

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)