Skip to content

Commit

Permalink
Add test if type with numpy array serializes back to ndarray
Browse files Browse the repository at this point in the history
  • Loading branch information
ketozhang committed Jun 16, 2024
1 parent 6a09af4 commit f91c69c
Showing 1 changed file with 48 additions and 0 deletions.
48 changes: 48 additions & 0 deletions tests/patterns/numpy_type_test.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,48 @@
import asdf
import numpy as np
from asdf.extension import Extension
from numpy.typing import NDArray

from asdf_pydantic import AsdfPydanticConverter, AsdfPydanticModel


class ArrayContainer(AsdfPydanticModel):
_tag = "asdf://asdf-pydantic/examples/tags/array-container-1.0.0"

array: NDArray # Equivalently np.ndarray


def setup_module():
"""Register the ArrayContainer model with the AsdfPydanticConverter.
Pytest will run this function before the tests in this module.
"""
AsdfPydanticConverter.add_models(ArrayContainer)

class TestExtension(Extension):
extension_uri = "asdf://asdf-pydantic/examples/extensions/test-1.0.0"

converters = [AsdfPydanticConverter()] # type: ignore
tags = [*AsdfPydanticConverter().tags] # type: ignore

asdf.get_config().add_extension(TestExtension())


########################################################################################
# Test Cases
########################################################################################


def test_convert_ArrayContainer_to_asdf(tmp_path):
"""When writing ArrayContainer to an ASDF file, the array field should be
serialized to the original numpy array.
"""
af = asdf.AsdfFile({"data": ArrayContainer(array=np.array([1, 2, 3]))}).write_to(
tmp_path / "test.asdf"
)

with asdf.open(tmp_path / "test.asdf") as af:
assert isinstance(af.tree["array"], np.ndarray), (
f"Expected {type(np.ndarray)}, " f"got {type(af.tree['array'])}"
)
assert np.all(af.tree["array"] == np.array([1, 2, 3]))

0 comments on commit f91c69c

Please sign in to comment.