diff --git a/dargs/_context.py b/dargs/_context.py new file mode 100644 index 0000000..16e3f50 --- /dev/null +++ b/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, + ) diff --git a/dargs/_refs.py b/dargs/_refs.py new file mode 100644 index 0000000..0f9fa7b --- /dev/null +++ b/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, + ) diff --git a/dargs/dargs.py b/dargs/dargs.py index 6a880e8..885a95d 100644 --- a/dargs/dargs.py +++ b/dargs/dargs.py @@ -21,7 +21,6 @@ import difflib import fnmatch import json -import os import re from copy import deepcopy from enum import Enum @@ -35,6 +34,10 @@ import typeguard +from ._context import TraversalContext +from ._refs import load_ref as _load_ref_file +from ._refs import resolve_ref + INDENT = " " # doc is indented by four spaces RAW_ANCHOR = False # whether to use raw html anchors or RST ones @@ -352,8 +355,38 @@ def traverse( _ref_base_dir: str | None = None, _trim_pattern: str | None = None, ) -> None: - # first, do something with the key - # then, take out the vaule and do something with it + """Traverse a mapping while applying the supplied hooks. + + The historical arguments are retained for callers that use this + low-level method directly. Internal recursion uses one context object + so reference and trimming state cannot be dropped between nodes. + """ + context = TraversalContext( + allow_ref=allow_ref, + trim_pattern=_trim_pattern, + ref_base_dir=_ref_base_dir, + ) + self._traverse( + argdict, + key_hook, + value_hook, + sub_hook, + variant_hook, + path, + context, + ) + + def _traverse( + self, + argdict: dict, + key_hook: HookArgKType, + value_hook: HookArgVType, + sub_hook: HookArgKType, + variant_hook: HookVrntType, + path: list[str] | None, + context: TraversalContext, + ) -> None: + """Traverse a mapping using an already initialized context.""" if path is None: path = [] key_hook(self, argdict, path) @@ -361,17 +394,14 @@ def traverse( value = argdict[self.name] value_hook(self, value, path) newpath = [*path, self.name] - # this is the key step that we traverse into the tree - self.traverse_value( + self._traverse_value( value, key_hook, value_hook, sub_hook, variant_hook, newpath, - allow_ref, - _ref_base_dir, - _trim_pattern, + context, ) def traverse_value( @@ -386,8 +416,33 @@ def traverse_value( _ref_base_dir: str | None = None, _trim_pattern: str | None = None, ) -> None: - # this is not private, and can be called directly - # in the condition where there is no leading key + """Traverse a value while applying the supplied hooks.""" + context = TraversalContext( + allow_ref=allow_ref, + trim_pattern=_trim_pattern, + ref_base_dir=_ref_base_dir, + ) + self._traverse_value( + value, + key_hook, + value_hook, + sub_hook, + variant_hook, + path, + context, + ) + + def _traverse_value( + self, + value: Any, + key_hook: HookArgKType, + value_hook: HookArgVType, + sub_hook: HookArgKType, + variant_hook: HookVrntType, + path: list[str] | None, + context: TraversalContext, + ) -> None: + """Traverse a value using an already initialized context.""" if path is None: path = [] if not self.repeat and isinstance(value, dict): @@ -398,9 +453,7 @@ def traverse_value( sub_hook, variant_hook, path, - allow_ref, - _ref_base_dir, - _trim_pattern, + context, ) elif self.repeat and isinstance(value, list): for idx, item in enumerate(value): @@ -411,16 +464,14 @@ def traverse_value( sub_hook, variant_hook, [*path, str(idx)], - allow_ref, - _ref_base_dir, - _trim_pattern, + context, ) elif self.repeat and isinstance(value, dict): # Repeat dictionaries use their keys as item names. Trim comment or # metadata entries before visiting items, since those entries may # not contain dictionaries and must not be type-checked as items. - if _trim_pattern is not None: - trim_by_pattern(value, _trim_pattern) + if context.trim_pattern is not None: + trim_by_pattern(value, context.trim_pattern) for kk, item in value.items(): self._traverse_sub( item, @@ -429,9 +480,7 @@ def traverse_value( sub_hook, variant_hook, [*path, kk], - allow_ref, - _ref_base_dir, - _trim_pattern, + context, ) def _traverse_sub( @@ -442,12 +491,12 @@ def _traverse_sub( sub_hook: HookArgKType = _DUMMYHOOK, variant_hook: HookVrntType = _DUMMYHOOK, path: list[str] | None = None, - allow_ref: bool = False, - _ref_base_dir: str | None = None, - _trim_pattern: str | None = None, + context: TraversalContext | None = None, ) -> None: if path is None: path = [self.name] + if context is None: + context = TraversalContext() if not isinstance(value, dict): raise ArgumentTypeError( path, @@ -456,22 +505,37 @@ def _traverse_sub( ) # A referenced file becomes the containing source for any nested refs # reached during this traversal. - ref_base_dir = _resolve_ref(value, allow_ref, _ref_base_dir) + ref_context = resolve_ref(value, context) sub_hook(self, value, path) for subvrnt in self.sub_variants.values(): variant_hook(subvrnt, value, path) for subarg in self.flatten_sub(value, path).values(): - subarg.traverse( - value, - key_hook, - value_hook, - sub_hook, - variant_hook, - path, - allow_ref, - ref_base_dir, - _trim_pattern, - ) + # Keep the historical dynamic dispatch for subclasses that + # override ``traverse``. Built-in Arguments use the private helper + # so the complete context (including reference provenance) can be + # threaded without reconstructing it from legacy parameters. + if type(subarg).traverse is Argument.traverse: + subarg._traverse( + value, + key_hook, + value_hook, + sub_hook, + variant_hook, + path, + ref_context, + ) + else: + subarg.traverse( + value, + key_hook, + value_hook, + sub_hook, + variant_hook, + path, + ref_context.allow_ref, + ref_context.ref_base_dir, + ref_context.trim_pattern, + ) # above are general traverse part # below are type checking part @@ -536,23 +600,24 @@ def check_value( A deep copy of ``value`` is made internally so the caller's data is not mutated. """ - ref_base_dir = None + context = TraversalContext(allow_ref=allow_ref) if allow_ref: value = deepcopy(value) # Resolve a root reference before validating its type or running # its extra check; traversal only resolves descendants. if isinstance(value, dict): - ref_base_dir = _resolve_ref(value, allow_ref) + context = resolve_ref(value, context) # ``traverse_value`` only checks descendants, so validate the root value # explicitly before descending into any sub-fields or variants. self._check_data(value, []) - self.traverse_value( + self._traverse_value( value, key_hook=Argument._check_exist, value_hook=Argument._check_data, sub_hook=Argument._check_strict if strict else _DUMMYHOOK, - allow_ref=allow_ref, - _ref_base_dir=ref_base_dir, + variant_hook=_DUMMYHOOK, + path=None, + context=context, ) def _check_exist(self, argdict: dict, path: list[str] | None = None) -> None: @@ -1243,116 +1308,32 @@ def trim_by_pattern( def _load_ref(ref_path: str) -> dict: - """Load a dict from an external file referenced by ``$ref``. - - Parameters - ---------- - ref_path : str - Path to the external file. Supported extensions: ``.json``, ``.yml``, ``.yaml``. + """Compatibility wrapper for the internal reference loader. Returns ------- dict - The loaded dict from the external file. - - Raises - ------ - ValueError - If the file extension is not supported, or if the file does not contain a - top-level mapping/object. - ImportError - If pyyaml is not installed and a YAML file is requested. + The mapping loaded from ``ref_path``. """ - 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 + return _load_ref_file(ref_path) def _resolve_ref(d: dict, allow_ref: bool = False, base_dir: str | None = None) -> str: - """Resolve the ``$ref`` key in a dict by loading from an external file. - - If ``$ref`` is present in ``d``, its value is treated as a file path. - The file is loaded and its contents are merged into ``d``. Keys already - present in ``d`` (other than ``$ref``) take precedence over keys from the - loaded file, allowing local overrides. Chained ``$ref`` values in the - loaded content are resolved in turn. Relative paths in a chain are resolved - from the directory of the file that contains them. Cyclic references are - detected and raise a ``ValueError``. + """Compatibility wrapper for the context-aware reference resolver. - The dict is modified **in place**. - - Parameters - ---------- - d : dict - The dict that may contain a ``$ref`` key. - allow_ref : bool, optional - If False (the default), raise a ``ValueError`` when ``$ref`` is found. - Set to True to enable loading from external files. - base_dir : str, optional - Directory containing ``d``. Relative references are resolved from this - directory; the process working directory is used when it is omitted. + The historical signature is retained for callers that import this private + helper. Traversal code uses :func:`resolve_ref` directly so the complete + reference chain remains available to child nodes. Returns ------- str The directory that nested mappings should use for relative references. - Raises - ------ - ValueError - If ``$ref`` is found but ``allow_ref`` is False, or if a cyclic - reference is detected. """ - if base_dir is None: - base_dir = os.curdir - if "$ref" not in d: - return base_dir - if not allow_ref: - raise ValueError( - "$ref is not allowed by default. " - "Pass allow_ref=True to enable loading from external files." - ) - visited_refs: set[str] = set() - while "$ref" in d: - ref_path = d.pop("$ref") - 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 visited_refs: - raise ValueError(f"Cyclic $ref detected for path: {canonical_ref_path!r}") - visited_refs.add(canonical_ref_path) - loaded = _load_ref(canonical_ref_path) - # A chained relative reference belongs to the file that declared it, - # rather than to the process's current working directory. - base_dir = os.path.dirname(canonical_ref_path) - # Merge: loaded content as base, local keys take precedence - merged = {**loaded, **d} - d.clear() - d.update(merged) - return base_dir + context = TraversalContext(allow_ref=allow_ref, ref_base_dir=base_dir) + resolved = resolve_ref(d, context) + return resolved.ref_base_dir or "." def isinstance_annotation(value: Any, dtype: type | Any) -> bool: diff --git a/dargs/notebook.py b/dargs/notebook.py index c4c836c..63e9377 100644 --- a/dargs/notebook.py +++ b/dargs/notebook.py @@ -27,7 +27,8 @@ from IPython.display import HTML, display from dargs import Argument, Variant -from dargs.dargs import _resolve_ref +from dargs._context import TraversalContext +from dargs._refs import resolve_ref __all__ = ["JSON"] @@ -174,6 +175,7 @@ def __init__( repeat: bool = False, allow_ref: bool = False, _ref_base_dir: str | None = None, + _ref_context: TraversalContext | None = None, ) -> None: self.data = data self.arg = arg @@ -181,7 +183,11 @@ def __init__( self.allow_ref = allow_ref # Keep the directory of the file that supplied this mapping so that # nested relative references are resolved beside their declaring file. - self._ref_base_dir = _ref_base_dir + self._ref_context = _ref_context or TraversalContext( + allow_ref=allow_ref, + ref_base_dir=_ref_base_dir, + ) + self._ref_base_dir = self._ref_context.ref_base_dir self.subdata = [] self._init_subdata() @@ -194,7 +200,7 @@ def _init_subdata(self) -> None: ): # Work on a copy to avoid mutating the caller's data data = self.data.copy() - ref_base_dir = _resolve_ref(data, self.allow_ref, self._ref_base_dir) + ref_context = resolve_ref(data, self._ref_context) sub_fields = self.arg.sub_fields.copy() # extend subfiles with sub_variants for vv in self.arg.sub_variants.values(): @@ -209,7 +215,7 @@ def _init_subdata(self) -> None: data[kk], sub_fields[kk], allow_ref=self.allow_ref, - _ref_base_dir=ref_base_dir, + _ref_context=ref_context, ) ) elif kk in self.arg.sub_variants: @@ -218,7 +224,7 @@ def _init_subdata(self) -> None: data[kk], self.arg.sub_variants[kk], allow_ref=self.allow_ref, - _ref_base_dir=ref_base_dir, + _ref_context=ref_context, ) ) else: @@ -236,7 +242,7 @@ def _init_subdata(self) -> None: self.arg, repeat=True, allow_ref=self.allow_ref, - _ref_base_dir=self._ref_base_dir, + _ref_context=self._ref_context, ) ) elif ( @@ -252,7 +258,7 @@ def _init_subdata(self) -> None: self.arg, repeat=True, allow_ref=self.allow_ref, - _ref_base_dir=self._ref_base_dir, + _ref_context=self._ref_context, ) ) diff --git a/tests/test_checker.py b/tests/test_checker.py index 653c0b5..68fd1ff 100644 --- a/tests/test_checker.py +++ b/tests/test_checker.py @@ -1,7 +1,7 @@ from __future__ import annotations import unittest -from typing import List +from typing import Any, List from dargs import Argument, Variant from dargs.dargs import ArgumentKeyError, ArgumentTypeError, ArgumentValueError @@ -100,6 +100,23 @@ def test_sub_fields(self) -> None: with self.assertRaises(ValueError): Argument("base", dict, [Argument("sub1", int), Argument("sub1", int)]) + def test_subclass_traverse_override_is_dispatched(self) -> None: + """Recursive traversal keeps honoring public ``traverse`` overrides.""" + + class TrackingArgument(Argument): + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + self.traverse_calls = 0 + + def traverse(self, *args: Any, **kwargs: Any) -> None: + self.traverse_calls += 1 + super().traverse(*args, **kwargs) + + child = TrackingArgument("child", dict, [Argument("value", int)]) + root = Argument("base", dict, [child]) + root.check({"base": {"child": {"value": 1}}}) + self.assertEqual(child.traverse_calls, 1) + def test_check_value_validates_root(self) -> None: """Root types and extra checks are enforced by check_value().""" with self.assertRaises(ArgumentTypeError): diff --git a/tests/test_ref.py b/tests/test_ref.py index c5a4db3..b77d12e 100644 --- a/tests/test_ref.py +++ b/tests/test_ref.py @@ -95,6 +95,28 @@ def test_ref_local_override(self) -> None: self.assertEqual(result["base"]["sub1"], 1) self.assertEqual(result["base"]["sub2"], "local") + def test_ref_repeated_reference_in_local_override(self) -> None: + """A local nested ref may independently reuse its parent's target.""" + shared_path = self._write_json("ref_repeated.json", {"value": 1}) + ca = Argument( + "base", + dict, + [ + Argument("value", int), + Argument("nested", dict, [Argument("value", int)]), + ], + ) + + ca.check( + { + "base": { + "$ref": shared_path, + "nested": {"$ref": shared_path}, + } + }, + allow_ref=True, + ) + def test_ref_yaml(self) -> None: """$ref to a YAML file is resolved when pyyaml is installed.""" if importlib.util.find_spec("yaml") is None: @@ -252,6 +274,37 @@ def test_ref_cyclic_detection(self) -> None: with self.assertRaises(ValueError, msg="Cyclic $ref"): ca.check({"base": {"$ref": ref_path}}, allow_ref=True) + def test_ref_nested_cycle_detection(self) -> None: + """Cycles crossing nested mappings are detected by the traversal context.""" + first_path = self._write_json( + "ref_nested_cycle_first.json", + {"nested": {"$ref": "ref_nested_cycle_second.json"}}, + ) + self._write_json( + "ref_nested_cycle_second.json", + {"nested": {"$ref": "ref_nested_cycle_first.json"}}, + ) + ca = Argument( + "base", + dict, + [ + Argument( + "nested", + dict, + [ + Argument( + "nested", + dict, + [Argument("nested", dict)], + ) + ], + ) + ], + ) + + with self.assertRaisesRegex(ValueError, "Cyclic \\$ref detected"): + ca.check({"base": {"$ref": first_path}}, allow_ref=True) + def test_ref_chained(self) -> None: """A nested relative $ref resolves beside the file that declares it.""" self._write_json("ref_inner.json", {"sub1": 7, "sub2": "inner"})