add load_yaml

asdf/_tests/test_util.py

@@ -1,8 +1,11 @@
+import contextlib
 import io
 import warnings
 
+import numpy as np
 import pytest
 
+import asdf
 from asdf import generic_io, util
 from asdf.exceptions import AsdfDeprecationWarning
 
@@ -116,3 +119,27 @@ def test_minversion():
 
         assert util.minversion(yaml, "3.1")
         assert util.minversion("yaml", "3.1")
+
+
+@pytest.mark.parametrize("input_type", ["filename", "binary_file", "generic_file"])
+@pytest.mark.parametrize("as_tagged", [True, False])
+def test_load_yaml(tmp_path, input_type, as_tagged):
+    fn = tmp_path / "test.asdf"
+    asdf.AsdfFile({"a": np.zeros(3)}).write_to(fn)
+
+    if input_type == "filename":
+        init = fn
+        ctx = contextlib.nullcontext()
+    elif input_type == "binary_file":
+        init = open(fn, "rb")
+        ctx = init
+    elif input_type == "generic_file":
+        init = generic_io.get_file(fn, "r")
+        ctx = init
+
+    with ctx:
+        tree = util.load_yaml(init, as_tagged=as_tagged)
+    if as_tagged:
+        assert isinstance(tree["a"], asdf.tagged.TaggedDict)
+    else:
+        assert not isinstance(tree["a"], asdf.tagged.TaggedDict)

asdf/util.py

@@ -11,6 +11,7 @@
 from urllib.request import pathname2url
 
 import numpy as np
+import yaml
 
 # The standard library importlib.metadata returns duplicate entrypoints
 # for all python versions up to and including 3.11
@@ -40,6 +41,7 @@
 
 
 __all__ = [
+    "load_yaml",
     "human_list",
     "get_array_base",
     "get_base_uri",
@@ -56,6 +58,50 @@
 ]
 
 
+def load_yaml(init, as_tagged=False):
+    """
+    Load just the yaml portion of an ASDF file
+
+    Parameters
+    ----------
+
+    init : filename or file-like
+        If file-like this must be opened in binary mode.
+
+    as_tagged: bool, optional
+        Return tree with instances of `asdf.tagged.Tagged` this
+        can be helpful if the yaml tags are of interest.
+        If False, the tree will only contain basic python types
+        (see the pyyaml ``BaseLoader`` documentation).
+
+    Returns
+    -------
+
+    tree : dict
+        Dictionary representing the ASDF tree
+    """
+
+    from .generic_io import get_file
+    from .yamlutil import AsdfLoader
+
+    if as_tagged:
+        loader = AsdfLoader
+    else:
+        loader = yaml.CBaseLoader if getattr(yaml, "__with_libyaml__", None) else yaml.BaseLoader
+
+    with get_file(init, "r") as gf:
+        reader = gf.reader_until(
+            constants.YAML_END_MARKER_REGEX,
+            7,
+            "End of YAML marker",
+            include=True,
+        )
+        # The following call to yaml.load is safe because we're
+        # using only loaders that don't create custom python objects
+        content = yaml.load(reader, Loader=loader)  # noqa: S506
+    return content
+
+
 def human_list(line, separator="and"):
     """
     Formats a list for human readability.