Skip to content

Commit f91c69c

Browse files
committed
Add test if type with numpy array serializes back to ndarray
1 parent 6a09af4 commit f91c69c

File tree

1 file changed

+48
-0
lines changed

1 file changed

+48
-0
lines changed

tests/patterns/numpy_type_test.py

+48
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,48 @@
1+
import asdf
2+
import numpy as np
3+
from asdf.extension import Extension
4+
from numpy.typing import NDArray
5+
6+
from asdf_pydantic import AsdfPydanticConverter, AsdfPydanticModel
7+
8+
9+
class ArrayContainer(AsdfPydanticModel):
10+
_tag = "asdf://asdf-pydantic/examples/tags/array-container-1.0.0"
11+
12+
array: NDArray # Equivalently np.ndarray
13+
14+
15+
def setup_module():
16+
"""Register the ArrayContainer model with the AsdfPydanticConverter.
17+
18+
Pytest will run this function before the tests in this module.
19+
"""
20+
AsdfPydanticConverter.add_models(ArrayContainer)
21+
22+
class TestExtension(Extension):
23+
extension_uri = "asdf://asdf-pydantic/examples/extensions/test-1.0.0"
24+
25+
converters = [AsdfPydanticConverter()] # type: ignore
26+
tags = [*AsdfPydanticConverter().tags] # type: ignore
27+
28+
asdf.get_config().add_extension(TestExtension())
29+
30+
31+
########################################################################################
32+
# Test Cases
33+
########################################################################################
34+
35+
36+
def test_convert_ArrayContainer_to_asdf(tmp_path):
37+
"""When writing ArrayContainer to an ASDF file, the array field should be
38+
serialized to the original numpy array.
39+
"""
40+
af = asdf.AsdfFile({"data": ArrayContainer(array=np.array([1, 2, 3]))}).write_to(
41+
tmp_path / "test.asdf"
42+
)
43+
44+
with asdf.open(tmp_path / "test.asdf") as af:
45+
assert isinstance(af.tree["array"], np.ndarray), (
46+
f"Expected {type(np.ndarray)}, " f"got {type(af.tree['array'])}"
47+
)
48+
assert np.all(af.tree["array"] == np.array([1, 2, 3]))

0 commit comments

Comments
 (0)