✨ BaseOutput: Add SubMapping for nested output namespaces

A code's output namespace can grow to dozens or even hundreds of fields once every
quantity it produces is exposed. Flattening all of them onto a single top-level mapping
floods autocomplete, makes the API hard to skim, and forces unrelated quantities
(forces, bands, occupations, magnetization, ...) to share one namespace with no
grouping. Related fields belong together under a named sub-namespace so users can
discover them by category instead of scrolling through one giant list.

`SubMapping` is a sentinel default, paired with the existing `@output_mapping`
decorator, that lets a mapping class declare a nested mapping as a regular
field:

    @output_mapping
    class _PwMapping:
        magnetization: _MagnetizationMapping

The decorator walks `get_type_hints` and injects a `SubMapping(mapping_cls=...)`
default for any bare annotation whose type is itself an `@output_mapping` class.

At runtime, `BaseOutput.outputs` recurses into each `SubMapping` field and builds the
nested mapping from the matching slice of the parsed output dict, so
`out.outputs.magnetization.total` works exactly like a top-level field — same lazy
resolution, same `AttributeError` on missing data, same frozen-dataclass features.

The constructor of `BaseOutput` now creates a (potentially) 2 level dictionary that
still maps the output names to their corresponding `Spec` and stores that in
`_output_spec_mapping`. The guard against non-`Spec` fields is still present for the
top level and sub mapping.

The `get_output` method is also updated. Currently, only top-level outputs are possible,
and these return a dictionary with all the available outputs in the sub mapping.
Conversion is also still possible for nested outputs via dot-separation, e.g.
to convert the "total" sub output in "magnetization" you have to define a converter for
"magnetization.total".

Author: mbercx

src/qe_tools/outputs/base.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import abc
+import contextlib
 import dataclasses
 import typing
 from functools import cached_property
@@ -16,34 +17,73 @@
 T = typing.TypeVar("T")
 
 
+class SubMapping:
+    """Sentinel marking a field as a nested output mapping.
+
+    `BaseOutput` resolves these at instantiation time. Nesting is intended to
+    be one level only: a sub-mapping class should only contain `Spec` fields.
+    """
+
+    def __init__(self, mapping_cls: type):
+        self.mapping_cls = mapping_cls
+
+
 def output_mapping(cls):
     """Decorator that defines a typed, frozen output mapping for a Quantum ESPRESSO code.
 
     Applies `@dataclass(frozen=True)` and injects `__getattribute__` and `__dir__` so that:
 
-    - Accessing a field whose value is still a `Spec` raises `AttributeError` with a clear
-      message (i.e. the output was not parsed).
+    - Accessing a field whose value is still a `Spec` or `SubMapping` raises
+      `AttributeError` with a clear message (i.e. the output was not parsed).
     - `dir()` only lists fields that were successfully extracted.
 
-    Each field must declare a `Spec(...)` as its default value:
+    Output fields declare a `Spec(...)` default. Sub-namespace fields are
+    declared with a bare annotation whose type is another `@output_mapping`
+    class — the decorator auto-injects a `SubMapping(hint)` default:
 
         fermi_energy: float = Spec("path.to.fermi_energy")
         \"""Fermi energy in eV.\"""
+        magnetization: _MagnetizationMapping
+        \"""Nested magnetization outputs.\"""
     """
 
     def __getattribute__(self, name):
         value = object.__getattribute__(self, name)
-        if isinstance(value, Spec):
+        if isinstance(value, (Spec, SubMapping)):
             raise AttributeError(f"'{name}' is not available in the parsed outputs.")
         return value
 
     def __dir__(self):
         return [
-            name for name, value in self.__dict__.items() if not isinstance(value, Spec)
+            name
+            for name, value in self.__dict__.items()
+            if not isinstance(value, (Spec, SubMapping))
         ]
 
     cls.__getattribute__ = __getattribute__
     cls.__dir__ = __dir__
+
+    # Inject `SubMapping(hint)` defaults for bare annotations whose type is
+    # itself an `@output_mapping`-decorated class. Note: `get_type_hints`
+    # evaluates annotations, which is fine as long as the module does not use
+    # `from __future__ import annotations` together with `TYPE_CHECKING`-only
+    # sub-mapping imports — in that case the eval would raise `NameError` here.
+    # Sub-mapping classes must be defined *before* the parent that references
+    # them (Python enforces this anyway for non-future-annotations modules).
+    for name, hint in typing.get_type_hints(cls).items():
+        if hasattr(cls, name):  # already has a default
+            continue
+
+        if not (isinstance(hint, type) and getattr(hint, "_is_output_mapping", False)):
+            raise TypeError(
+                f"{cls.__name__}.{name}: needs a Spec(...) default, or a bare "
+                f"annotation whose type is an @output_mapping class "
+                f"(which must be defined before this class)"
+            )
+
+        setattr(cls, name, SubMapping(hint))
+
+    cls._is_output_mapping = True
     return dataclasses.dataclass(frozen=True)(cls)
 
 
@@ -68,15 +108,25 @@ def _get_mapping_class(cls) -> type:
 
     def __init__(self, raw_outputs: dict):
         self.raw_outputs = raw_outputs
-        self._output_spec_mapping = {}
 
-        for field in dataclasses.fields(self._get_mapping_class()):
-            if not isinstance(field.default, Spec):
-                raise TypeError(
-                    f"{type(self).__name__}.{field.name}: expected a Spec(...) default, "
-                    f"got {field.default!r}"
-                )
-            self._output_spec_mapping[field.name] = field.default
+        def build(mapping_cls: type) -> dict:
+            """Build the nested spec dict from a mapping class."""
+            result: dict = {}
+
+            for field in dataclasses.fields(mapping_cls):
+                if isinstance(field.default, SubMapping):
+                    result[field.name] = build(field.default.mapping_cls)
+                elif isinstance(field.default, Spec):
+                    result[field.name] = field.default
+                else:
+                    raise TypeError(
+                        f"{mapping_cls.__name__}.{field.name}: expected a Spec(...) or "
+                        f"SubMapping(...) default, got {field.default!r}"
+                    )
+
+            return result
+
+        self._output_spec_mapping = build(self._get_mapping_class())
 
     @classmethod
     @abc.abstractmethod
@@ -109,7 +159,16 @@ def get_output(
             >>> pw_out.get_output(name="structure")
             >>> pw_out.get_output(name="structure", to="pymatgen")
         """
-        output_data = glom(self.raw_outputs, self._output_spec_mapping[name])
+        entry = self._output_spec_mapping[name]
+
+        if isinstance(entry, dict):
+            output_data: typing.Any = {}
+
+            for sub_name, sub_spec in entry.items():
+                with contextlib.suppress(GlomError):
+                    output_data[sub_name] = glom(self.raw_outputs, sub_spec)
+        else:
+            output_data = glom(self.raw_outputs, entry)
 
         if to is None:
             return output_data
@@ -126,6 +185,14 @@ def get_output(
         else:
             raise ValueError(f"Library '{to}' is not supported.")
 
+        if isinstance(entry, dict):
+            return {
+                sub_name: Converter().convert(f"{name}.{sub_name}", sub_value)
+                if sub_name in Converter.conversion_mapping
+                else sub_value
+                for sub_name, sub_value in output_data.items()
+            }
+
         return (
             Converter().convert(name, output_data)
             if name in Converter.conversion_mapping
@@ -180,4 +247,15 @@ def list_outputs(self, only_available: bool = True) -> list[str]:
     @cached_property
     def outputs(self) -> T:
         """Namespace with available outputs."""
-        return self._get_mapping_class()(**self.get_output_dict())
+
+        def build(mapping_cls: type, data: dict):
+            defaults = {f.name: f.default for f in dataclasses.fields(mapping_cls)}
+            kwargs = {
+                name: build(defaults[name].mapping_cls, value)  # type: ignore[union-attr]
+                if isinstance(defaults[name], SubMapping)
+                else value
+                for name, value in data.items()
+            }
+            return mapping_cls(**kwargs)
+
+        return build(self._get_mapping_class(), self.get_output_dict())

tests/outputs/test_base.py

@@ -78,3 +78,89 @@ def test_get_output_dict(raw_outputs):
                 "B",
             ]
         )
+
+
+# --- SubMapping (nested output namespaces) ----------------------------------
+
+
+@output_mapping
+class _NestedMapping:
+    c: int = Spec("b.c")
+    d: int = Spec("b.d")
+    missing: int = Spec("b.nope")
+
+
+@output_mapping
+class _ParentMapping:
+    A: float = Spec("a")
+    nested: _NestedMapping
+
+
+class _ParentOutput(BaseOutput[_ParentMapping]):
+    @classmethod
+    def from_dir(cls, _: str):
+        pass
+
+
+def test_submapping_output_access(raw_outputs):
+    """Resolved outputs on a sub-namespace are accessible via attribute."""
+    outputs = _ParentOutput(raw_outputs).outputs
+    assert outputs.nested.c == 3
+    assert outputs.nested.d == 4
+
+
+def test_submapping_missing_output_raises(raw_outputs):
+    outputs = _ParentOutput(raw_outputs).outputs
+    with pytest.raises(AttributeError, match="missing.*not available"):
+        outputs.nested.missing
+
+
+def test_submapping_list_outputs_top_level_only(raw_outputs):
+    """`list_outputs` yields top-level field names only — sub-namespaces as a single entry."""
+    pw_out = _ParentOutput(raw_outputs)
+    # Both modes return the same result here: sub-namespaces are *always* listed
+    # (even if every output is missing), and there are no unresolvable top-level
+    # outputs in `_ParentMapping` for `only_available=True` to filter out.
+    assert pw_out.list_outputs() == ["A", "nested"]
+    assert pw_out.list_outputs(only_available=False) == ["A", "nested"]
+
+
+def test_submapping_get_output_namespace_returns_dict(raw_outputs):
+    """`get_output(<sub-namespace>)` returns a partial dict of available outputs."""
+    pw_out = _ParentOutput(raw_outputs)
+    assert pw_out.get_output("nested") == {"c": 3, "d": 4}
+    # Users index the dict directly.
+    assert pw_out.get_output("nested")["c"] == 3
+
+
+def test_submapping_get_output_dict_shape(raw_outputs):
+    """`get_output_dict()` is flat at the top level; sub-namespaces are nested dicts."""
+    assert _ParentOutput(raw_outputs).get_output_dict() == {
+        "A": 1,
+        "nested": {"c": 3, "d": 4},
+    }
+
+
+def test_validation_rejects_non_spec_top_level_default():
+    """Top-level field defaults that aren't `Spec`/`SubMapping` raise in `__init__`."""
+
+    @output_mapping
+    class _BadParent:
+        bad: int = 42  # type: ignore[assignment]
+
+    class _BadOutput(BaseOutput[_BadParent]):
+        @classmethod
+        def from_dir(cls, _: str):
+            pass
+
+    with pytest.raises(TypeError, match="_BadParent.bad"):
+        _BadOutput(raw_outputs={})
+
+
+def test_decorator_rejects_bare_annotation_non_output_mapping():
+    """Bare annotation whose type isn't `@output_mapping`-decorated is rejected at decoration time."""
+    with pytest.raises(TypeError, match="bad.*@output_mapping class"):
+
+        @output_mapping
+        class _BadBare:
+            bad: int