add read_buffer_context and write_read_buffer_context (#875)

* add read_buffer_context and write_read_buffer_context closes #873 * update changelog * update dummy test * add read_buffer_context and write_read_buffer_context closes #873 * update changelog * update dummy test * removed internal keyword dummy_inline_arrays --------- Co-authored-by: Çağtay Fabry <43667554+CagtayFabry@users.noreply.github.com> Co-authored-by: Martin K. Scherer <m.scherer@fu-berlin.de>
BAMWelDX · May 16, 2023 · e4e563c · e4e563c
1 parent 7b5c0c1
commit e4e563c
Show file tree

Hide file tree

Showing 13 changed files with 243 additions and 194 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -4,7 +4,14 @@
 
 ### Fixes
 
-- fix compatibility with `pint=0.21` \[{pull}`876`\] .
+- fix compatibility with `pint=0.21` \[{pull}`876`\].
+- add `read_buffer_context` and `write_read_buffer_context` to `weldx.asdf.util`
+  to fix tests accessing closed files. {issue}`873` \[{pull}`875`\].
+
+### Changes
+
+- removed keyword `dummy_inline_arrays` from function `asdf.util.write_buffer`
+  as it was only used internally \[{pull}`875`\].
 
 ## 0.6.6 (19.04.2023)
 

diff --git a/weldx/asdf/util.py b/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",
@@ -96,19 +99,12 @@ def write_buffer(
     -------
     io.BytesIO
         Bytes buffer of the ASDF file.
-
-    Notes
-    -----
-    In addition to the usual asdf.AsdfFile.write_to arguments in write_args you can pass
-    the parameter "dummy_arrays". If set, all array data is replaced with a empty list.
     """
     if asdffile_kwargs is None:
         asdffile_kwargs = {}
     if write_kwargs is None:
         write_kwargs = {}
 
-    dummy_inline_arrays = write_kwargs.pop("dummy_arrays", False)
-
     if _use_weldx_file is None:
         _use_weldx_file = _USE_WELDX_FILE
 
@@ -120,7 +116,6 @@ def show(wx):
             wx.header(False, False)
 
     if _use_weldx_file:
-        write_kwargs = dict(all_array_storage="inline")
         from weldx.asdf.file import WeldxFile
 
         with WeldxFile(
@@ -129,30 +124,23 @@ def show(wx):
             write_kwargs=write_kwargs,
             mode="rw",
         ) as wx:
-            wx.write_to()
             show(wx)
             buff = wx.file_handle
     else:
         buff = BytesIO()
         with asdf.AsdfFile(tree, extensions=None, **asdffile_kwargs) as ff:
-            if dummy_inline_arrays:  # lets store an empty list in the asdf file.
-                write_kwargs["all_array_storage"] = "inline"
-                from unittest.mock import patch
-
-                with patch("asdf.tags.core.ndarray.numpy_array_to_list", lambda x: []):
-                    ff.write_to(buff, **write_kwargs)
-            else:
-                ff.write_to(buff, **write_kwargs)
+            ff.write_to(buff, **write_kwargs)
     buff.seek(0)
     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 +157,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 +173,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 +254,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]:

diff --git a/weldx/tests/asdf_tests/test_asdf_aws_schema.py b/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)
diff --git a/weldx/tests/asdf_tests/test_asdf_base_types.py b/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)