add read_buffer_context and write_read_buffer_context

closes BAMWelDX#873

weldx/asdf/util.py

@@ -2,6 +2,7 @@
 from __future__ import annotations
 
 from collections.abc import Callable, Mapping, Set
+from contextlib import contextmanager
 from io import BytesIO, TextIOBase
 from pathlib import Path
 from typing import Any, Hashable, MutableMapping, Union
@@ -34,8 +35,10 @@
 __all__ = [
     "get_schema_path",
     "read_buffer",
+    "read_buffer_context",
     "write_buffer",
     "write_read_buffer",
+    "write_read_buffer_context",
     "get_yaml_header",
     "view_tree",
     "notebook_fileprinter",
@@ -147,12 +150,13 @@ def show(wx):
     return buff
 
 
-def read_buffer(
+@contextmanager
+def read_buffer_context(
     buffer: BytesIO,
     open_kwargs: dict = None,
     _use_weldx_file=_USE_WELDX_FILE,
 ):
-    """Read ASDF file contents from buffer instance.
+    """Contextmanager to read ASDF file contents from buffer instance.
 
     Parameters
     ----------
@@ -169,7 +173,7 @@ def read_buffer(
 
     """
     if open_kwargs is None:
-        open_kwargs = {"copy_arrays": True}
+        open_kwargs = {"copy_arrays": True, "lazy_load": False}
 
     buffer.seek(0)
 
@@ -185,8 +189,62 @@ def read_buffer(
         extensions=None,
         **open_kwargs,
     ) as af:
-        data = af.tree
-    return data
+        yield af.tree
+
+
+def read_buffer(
+    buffer: BytesIO,
+    open_kwargs: dict = None,
+    _use_weldx_file=_USE_WELDX_FILE,
+):
+    """Read ASDF file contents from buffer instance.
+
+    Parameters
+    ----------
+    buffer : io.BytesIO
+        Buffer containing ASDF file contents
+    open_kwargs
+        Additional keywords to pass to `asdf.AsdfFile.open`
+        Extensions are always set, ``copy_arrays=True`` is set by default.
+
+    Returns
+    -------
+    dict
+        ASDF file tree.
+
+    """
+    with read_buffer_context(buffer, open_kwargs, _use_weldx_file) as data:
+        return data
+
+
+@contextmanager
+def write_read_buffer_context(
+    tree: dict, asdffile_kwargs=None, write_kwargs=None, open_kwargs=None
+):
+    """Context manager to perform a buffered write/read roundtrip of a tree
+    using default ASDF settings.
+
+    Parameters
+    ----------
+    tree
+        Tree object to serialize.
+    asdffile_kwargs
+        Additional keywords to pass to `asdf.AsdfFile`
+    write_kwargs
+        Additional keywords to pass to `asdf.AsdfFile.write_to`
+        Extensions are always set.
+    open_kwargs
+        Additional keywords to pass to `asdf.AsdfFile.open`
+        Extensions are always set, ``copy_arrays=True`` is set by default.
+
+    Returns
+    -------
+    dict
+
+    """
+    buffer = write_buffer(tree, asdffile_kwargs, write_kwargs)
+    with read_buffer_context(buffer, open_kwargs) as data:
+        yield data
 
 
 def write_read_buffer(
@@ -212,8 +270,11 @@ def write_read_buffer(
     dict
 
     """
-    buffer = write_buffer(tree, asdffile_kwargs, write_kwargs)
-    return read_buffer(buffer, open_kwargs)
+
+    with write_read_buffer_context(
+        tree, asdffile_kwargs, write_kwargs, open_kwargs
+    ) as data:
+        return data
 
 
 def get_yaml_header(file: types_path_and_file_like, parse=False) -> Union[str, dict]:

weldx/tests/asdf_tests/test_asdf_aws_schema.py

@@ -1,7 +1,7 @@
 """Test ASDF serialization of AWS schema definitions."""
 import pytest
 
-from weldx.asdf.util import write_read_buffer
+from weldx.asdf.util import write_read_buffer_context
 from weldx.constants import Q_
 
 # weld design -----------------------------------------------------------------
@@ -100,7 +100,7 @@ def test_aws_example():
 
     tree = dict(process=process, weldment=weldment, base_metal=base_metal)
 
-    data = write_read_buffer(tree)
-    data.pop("asdf_library", None)
-    data.pop("history", None)
-    assert compare_nested(data, tree)
+    with write_read_buffer_context(tree) as data:
+        data.pop("asdf_library", None)
+        data.pop("history", None)
+        assert compare_nested(data, tree)

weldx/tests/asdf_tests/test_asdf_base_types.py

@@ -6,7 +6,7 @@
 import xarray as xr
 from asdf import ValidationError
 
-from weldx.asdf.util import write_read_buffer
+from weldx.asdf.util import write_read_buffer, write_read_buffer_context
 
 
 # --------------------------------------------------------------------------------------
@@ -28,12 +28,14 @@ def test_dataarray():
     da.attrs["long_name"] = "random velocity"
     # add metadata to coordinate
     da.x.attrs["units"] = "x units"
-    da2 = write_read_buffer({"da": da})["da"]
-    assert da2.identical(da)
+    with write_read_buffer_context({"da": da}) as data:
+        da2 = data["da"]
+        assert da2.identical(da)
 
 
 def test_dataset_children():
     da = xr.DataArray(np.arange(10), name="arr1", attrs={"name": "sample data"})
     ds = da.to_dataset()
-    ds2 = write_read_buffer({"ds": ds})["ds"]
-    assert ds2.identical(ds)
+    with write_read_buffer_context({"ds": ds}) as data:
+        ds2 = data["ds"]
+        assert ds2.identical(ds)