refactor: centralize traversal context and ref loading

dargs/_context.py

@@ -0,0 +1,95 @@
+"""Internal state shared by dargs tree traversals."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, replace
+
+
+@dataclass(frozen=True)
+class TraversalContext:
+    """Carry operation-scoped state while walking an argument tree.
+
+    Keeping path-sensitive options together prevents individual traversal
+    helpers from silently dropping state when they recurse into child
+    arguments. The reference chain is immutable so sibling branches do not
+    accidentally look like cyclic references to one another.
+    """
+
+    allow_ref: bool = False
+    trim_pattern: str | None = None
+    ref_base_dir: str | None = None
+    ref_chain: tuple[str, ...] = ()
+    # Mapping identities retain the source context of values merged from a
+    # reference.  This lets local overrides resolve independently while still
+    # detecting cycles through mappings that actually came from a referenced
+    # file.
+    ref_origins: tuple[tuple[int, str | None, tuple[str, ...]], ...] = ()
+
+    def for_mapping(self, mapping: object) -> TraversalContext:
+        """Return the context associated with ``mapping`` when known.
+
+        A single merged dictionary can contain values from several sources:
+        keys supplied locally and keys loaded from ``$ref``.  Traversal uses
+        this mapping-specific provenance to select the correct base directory
+        and active reference ancestry for each nested mapping.
+
+        Returns
+        -------
+        TraversalContext
+            A context using the mapping-specific reference state when known.
+        """
+        mapping_id = id(mapping)
+        for origin_id, base_dir, ref_chain in reversed(self.ref_origins):
+            if origin_id == mapping_id:
+                return replace(
+                    self,
+                    ref_base_dir=base_dir,
+                    ref_chain=ref_chain,
+                )
+        return self
+
+    def with_mapping_origins(
+        self,
+        origins: dict[int, tuple[str | None, tuple[str, ...]]],
+    ) -> TraversalContext:
+        """Return a context extended with mapping provenance entries.
+
+        Returns
+        -------
+        TraversalContext
+            A context containing the supplied provenance in addition to the
+            existing entries.
+        """
+        if not origins:
+            return self
+        merged = {
+            origin_id: (base_dir, ref_chain)
+            for origin_id, base_dir, ref_chain in self.ref_origins
+        }
+        merged.update(origins)
+        return replace(
+            self,
+            ref_origins=tuple(
+                (origin_id, base_dir, ref_chain)
+                for origin_id, (base_dir, ref_chain) in merged.items()
+            ),
+        )
+
+    def with_ref_state(
+        self,
+        *,
+        ref_base_dir: str,
+        ref_chain: tuple[str, ...],
+    ) -> TraversalContext:
+        """Return a child context with updated reference resolution state.
+
+        Returns
+        -------
+        TraversalContext
+            A context carrying the supplied reference state.
+        """
+        return replace(
+            self,
+            ref_base_dir=ref_base_dir,
+            ref_chain=ref_chain,
+        )

dargs/_refs.py

@@ -0,0 +1,168 @@
+"""Loading and resolving external ``$ref`` mappings."""
+
+from __future__ import annotations
+
+import json
+import os
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from collections.abc import Iterator
+
+    from ._context import TraversalContext
+
+__all__ = ["load_ref", "resolve_ref"]
+
+
+def _mapping_origins(
+    value: object,
+    origin: tuple[str | None, tuple[str, ...]],
+    seen: set[int] | None = None,
+) -> Iterator[tuple[int, tuple[str | None, tuple[str, ...]]]]:
+    """Yield provenance entries for mappings nested inside ``value``.
+
+    Yields
+    ------
+    tuple[int, tuple[str | None, tuple[str, ...]]]
+        A mapping identity and its source base directory/reference chain.
+    """
+    if seen is None:
+        seen = set()
+    if isinstance(value, dict):
+        value_id = id(value)
+        if value_id in seen:
+            return
+        seen.add(value_id)
+        yield value_id, origin
+        for child in value.values():
+            yield from _mapping_origins(child, origin, seen)
+    elif isinstance(value, list):
+        for child in value:
+            yield from _mapping_origins(child, origin, seen)
+
+
+def load_ref(ref_path: str) -> dict:
+    """Load a mapping from a JSON or YAML file referenced by ``$ref``.
+
+    Parameters
+    ----------
+    ref_path : str
+        Path to the external file. Supported extensions are ``.json``,
+        ``.yml``, and ``.yaml``.
+
+    Returns
+    -------
+    dict
+        The loaded mapping.
+
+    Raises
+    ------
+    ValueError
+        If the extension is unsupported or the file does not contain a
+        top-level mapping.
+    ImportError
+        If a YAML file is requested without PyYAML installed.
+    """
+    ext = os.path.splitext(ref_path)[1].lower()
+    if ext == ".json":
+        with open(ref_path, encoding="utf-8") as f:
+            loaded = json.load(f)
+    elif ext in (".yml", ".yaml"):
+        try:
+            import yaml
+        except ImportError as e:
+            raise ImportError(
+                "pyyaml is required to load YAML files referenced by $ref. "
+                "Install it with: pip install pyyaml"
+            ) from e
+        with open(ref_path, encoding="utf-8") as f:
+            loaded = yaml.safe_load(f)
+    else:
+        raise ValueError(
+            f"Unsupported file extension `{ext}` for $ref. "
+            "Supported extensions are: .json, .yml, .yaml"
+        )
+    if not isinstance(loaded, dict):
+        raise ValueError(
+            f"Referenced file {ref_path!r} must contain a mapping/object at the top "
+            f"level, but got {type(loaded).__name__!r}."
+        )
+    return loaded
+
+
+def resolve_ref(d: dict, context: TraversalContext) -> TraversalContext:
+    """Resolve ``$ref`` entries and return the child traversal context.
+
+    Relative references are resolved from the file that supplied the current
+    mapping. ``context.ref_chain`` tracks active ancestor files for mappings
+    loaded from references. Local overrides retain their original provenance,
+    so a finite repeated reference is not mistaken for a cycle.
+
+    The mapping is modified in place, matching the historical private
+    ``dargs.dargs._resolve_ref`` helper.
+
+    Returns
+    -------
+    TraversalContext
+        Context updated with the directory and active reference chain for
+        descendants.
+
+    Raises
+    ------
+    ValueError
+        If references are disabled or a cyclic reference is detected.
+    """
+    context = context.for_mapping(d)
+    base_dir = context.ref_base_dir if context.ref_base_dir is not None else os.curdir
+    if "$ref" not in d:
+        return context.with_ref_state(
+            ref_base_dir=base_dir,
+            ref_chain=context.ref_chain,
+        )
+    if not context.allow_ref:
+        raise ValueError(
+            "$ref is not allowed by default. "
+            "Pass allow_ref=True to enable loading from external files."
+        )
+
+    ref_chain = context.ref_chain
+    origins = {
+        origin_id: (origin_base_dir, origin_chain)
+        for origin_id, origin_base_dir, origin_chain in context.ref_origins
+    }
+    while "$ref" in d:
+        ref_path = d.pop("$ref")
+        # Values already present in ``d`` are local to the current source. A
+        # chained reference may merge another source on top, but local values
+        # must keep this state for their own nested references.
+        local_items = dict(d)
+        local_origin = (base_dir, ref_chain)
+        resolved_ref_path = (
+            ref_path if os.path.isabs(ref_path) else os.path.join(base_dir, ref_path)
+        )
+        canonical_ref_path = os.path.realpath(resolved_ref_path)
+        if canonical_ref_path in ref_chain:
+            raise ValueError(f"Cyclic $ref detected for path: {canonical_ref_path!r}")
+        ref_chain = (*ref_chain, canonical_ref_path)
+        loaded = load_ref(canonical_ref_path)
+        # A chained relative reference belongs to the file that declares it.
+        base_dir = os.path.dirname(canonical_ref_path)
+        loaded_origin = (base_dir, ref_chain)
+        # Preserve provenance on both sides of the merge. ``setdefault`` keeps
+        # values retained from an earlier source correctly labeled when a
+        # chained reference adds another layer.
+        for value in local_items.values():
+            for origin_id, origin in _mapping_origins(value, local_origin):
+                origins.setdefault(origin_id, origin)
+        for key, value in loaded.items():
+            if key not in local_items:
+                for origin_id, origin in _mapping_origins(value, loaded_origin):
+                    origins.setdefault(origin_id, origin)
+        merged = {**loaded, **local_items}
+        d.clear()
+        d.update(merged)
+
+    return context.with_mapping_origins(origins).with_ref_state(
+        ref_base_dir=base_dir,
+        ref_chain=ref_chain,
+    )