Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

add read_buffer_context and write_read_buffer_context #875

Merged
merged 10 commits into from
May 16, 2023
9 changes: 8 additions & 1 deletion CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand Down
93 changes: 69 additions & 24 deletions weldx/asdf/util.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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",
Expand Down Expand Up @@ -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

Expand All @@ -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(
Expand All @@ -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
----------
Expand All @@ -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)

Expand All @@ -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(
Expand All @@ -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]:
Expand Down
10 changes: 5 additions & 5 deletions weldx/tests/asdf_tests/test_asdf_aws_schema.py
Original file line number Diff line number Diff line change
@@ -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 -----------------------------------------------------------------
Expand Down Expand Up @@ -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)
12 changes: 7 additions & 5 deletions weldx/tests/asdf_tests/test_asdf_base_types.py
Original file line number Diff line number Diff line change
Expand Up @@ -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


# --------------------------------------------------------------------------------------
Expand All @@ -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)
Loading