From 4b1a625037673710b4058bf636bbd51b6dd4ec05 Mon Sep 17 00:00:00 2001 From: Brett Date: Fri, 8 Dec 2023 16:28:07 -0500 Subject: [PATCH] add tree_type yaml --- asdf/_asdf.py | 43 ++++++++++++++++------------------- asdf/_tests/_helpers.py | 8 +++---- asdf/_tests/test_asdf.py | 4 +++- asdf/_tests/test_reference.py | 4 +--- 4 files changed, 27 insertions(+), 32 deletions(-) diff --git a/asdf/_asdf.py b/asdf/_asdf.py index a5bd4a79f..a199f04f5 100644 --- a/asdf/_asdf.py +++ b/asdf/_asdf.py @@ -6,6 +6,7 @@ import time import warnings +import yaml from packaging.version import Version from . import _compression as mcompression @@ -28,6 +29,8 @@ from .tags.core import AsdfObject, ExtensionMetadata, HistoryEntry, Software from .util import NotSet +_yaml_base_loader = yaml.CBaseLoader if getattr(yaml, "__with_libyaml__", None) else yaml.BaseLoader + def __getattr__(name): if name == "SerializationContext": @@ -762,13 +765,22 @@ def _open_asdf( fd, validate_checksums=False, extensions=None, - tree_type="custom", - _get_yaml_content=False, + tree_type=None, _force_raw_types=NotSet, strict_extension_check=False, ignore_missing_extensions=False, ): """Attempt to populate AsdfFile data from file-like object""" + if tree_type is None: + tree_type = "custom" + + if _force_raw_types is not NotSet: + warnings.warn("_force_raw_types is deprecated and will be replaced by tree_type", AsdfDeprecationWarning) + tree_type = "tagged" if _force_raw_types else tree_type + + if tree_type not in ("custom", "tagged", "yaml"): + msg = f"Unsupported tree type {tree_type}" + raise ValueError(msg) if strict_extension_check and ignore_missing_extensions: msg = "'strict_extension_check' and 'ignore_missing_extensions' are incompatible options" @@ -815,11 +827,10 @@ def _open_asdf( initial_content=yaml_token, ) - # For testing: just return the raw YAML content - if _get_yaml_content: - yaml_content = reader.read() - fd.close() - return yaml_content + # just return the raw YAML content + if tree_type == "yaml": + self._tree = yaml.load(reader.read(), Loader=_yaml_base_loader) # noqa: S506 + return self # We parse the YAML content into basic data structures # now, but we don't do anything special with it until @@ -852,19 +863,6 @@ def _open_asdf( self.close() raise - if _force_raw_types is not NotSet: - warnings.warn( - "_force_raw_types is deprecated and will be replaced by tree_type", AsdfDeprecationWarning - ) - tree_type = "tagged" if _force_raw_types else "custom" - - if tree_type is None: - tree_type = "custom" - - if tree_type not in ("custom", "tagged"): - msg = f"Unsupported tree type {tree_type}" - raise ValueError(msg) - if tree_type == "custom": tree = yamlutil.tagged_tree_to_custom_tree(tree, self) @@ -885,7 +883,6 @@ def _open_impl( validate_checksums=False, extensions=None, tree_type=None, - _get_yaml_content=False, _force_raw_types=NotSet, strict_extension_check=False, ignore_missing_extensions=False, @@ -900,7 +897,6 @@ def _open_impl( validate_checksums=validate_checksums, extensions=extensions, tree_type=tree_type, - _get_yaml_content=_get_yaml_content, _force_raw_types=_force_raw_types, strict_extension_check=strict_extension_check, ignore_missing_extensions=ignore_missing_extensions, @@ -1511,7 +1507,6 @@ def open_asdf( custom_schema=None, strict_extension_check=False, ignore_missing_extensions=False, - _get_yaml_content=False, ): """ Open an existing ASDF file. @@ -1584,6 +1579,7 @@ def open_asdf( - "custom" (default) for a tree with custom objects - "tagged" for a tree with `asdf.tagged` objects + - "yaml" return the raw yaml contents of just file Returns ------- @@ -1620,7 +1616,6 @@ def open_asdf( validate_checksums=validate_checksums, extensions=extensions, tree_type=tree_type, - _get_yaml_content=_get_yaml_content, _force_raw_types=_force_raw_types, strict_extension_check=strict_extension_check, ignore_missing_extensions=ignore_missing_extensions, diff --git a/asdf/_tests/_helpers.py b/asdf/_tests/_helpers.py index 693378f68..2fb2d5c1b 100644 --- a/asdf/_tests/_helpers.py +++ b/asdf/_tests/_helpers.py @@ -25,7 +25,7 @@ import asdf from asdf import generic_io, versioning from asdf._asdf import AsdfFile, _get_asdf_library_info -from asdf.constants import YAML_TAG_PREFIX +from asdf.constants import YAML_END_MARKER_REGEX, YAML_TAG_PREFIX from asdf.exceptions import AsdfConversionWarning from asdf.tags.core import AsdfObject from asdf.versioning import ( @@ -219,13 +219,13 @@ def _assert_roundtrip_tree( asdf_check_func(ff) buff.seek(0) - ff = AsdfFile(extensions=extensions, **init_options) - content = AsdfFile._open_impl(ff, buff, mode="r", _get_yaml_content=True) + gf = generic_io.get_file(buff) + content = gf.reader_until(YAML_END_MARKER_REGEX, 7, "End of YAML marker", include=True).read() buff.close() # We *never* want to get any raw python objects out assert b"!!python" not in content assert b"!core/asdf" in content - assert content.startswith(b"%YAML 1.1") + assert b"%YAML 1.1" in content if raw_yaml_check_func: raw_yaml_check_func(content) diff --git a/asdf/_tests/test_asdf.py b/asdf/_tests/test_asdf.py index e4a079585..532d34311 100644 --- a/asdf/_tests/test_asdf.py +++ b/asdf/_tests/test_asdf.py @@ -394,7 +394,7 @@ def test_fsspec_http(httpserver): assert_tree_match(tree, af.tree) -@pytest.mark.parametrize("tree_type", [None, "custom", "tagged", "unknown"]) +@pytest.mark.parametrize("tree_type", [None, "custom", "tagged", "yaml", "unknown"]) def test_asdf_open_tree_type(tmp_path, tree_type): fn = tmp_path / "test.asdf" asdf.AsdfFile({"a": np.zeros(3)}).write_to(fn) @@ -407,5 +407,7 @@ def test_asdf_open_tree_type(tmp_path, tree_type): with asdf.open(fn, tree_type=tree_type) as af: if tree_type == "tagged": assert isinstance(af["a"], asdf.tagged.TaggedDict) + elif tree_type == "yaml": + assert isinstance(af["a"], dict) else: assert isinstance(af["a"], asdf.tags.core.ndarray.NDArrayType) diff --git a/asdf/_tests/test_reference.py b/asdf/_tests/test_reference.py index d9186fa08..957c8cb68 100644 --- a/asdf/_tests/test_reference.py +++ b/asdf/_tests/test_reference.py @@ -190,9 +190,7 @@ def test_internal_reference(tmp_path): buff = io.BytesIO() ff.write_to(buff) buff.seek(0) - ff = asdf.AsdfFile() - content = asdf.AsdfFile()._open_impl(ff, buff, _get_yaml_content=True) - assert b"{$ref: ''}" in content + assert b"{$ref: ''}" in buff.read() def test_implicit_internal_reference(tmp_path):