feat: Add before-validation converts all ASDF NDArrayType to ndarray

Author: ketozhang

asdf_pydantic/model.py

@@ -1,8 +1,11 @@
 import textwrap
-from typing import ClassVar
+from typing import Any, ClassVar
 
+import numpy as np
 import yaml
-from pydantic import BaseModel
+from asdf.tags.core import NDArrayType
+from numpy.typing import NDArray
+from pydantic import BaseModel, ValidationInfo, field_validator
 
 
 class AsdfPydanticModel(BaseModel):
@@ -67,3 +70,14 @@ def schema_asdf(
         )
         body = yaml.dump(cls.schema())
         return header + body
+
+    @field_validator("*", mode="before")
+    @classmethod
+    def _allow_asdf_NDArrayType_to_be_ndarray(
+        cls, value: Any, info: ValidationInfo
+    ) -> Any | NDArray:
+        """Before Pydantic validation, convert NDArrayType to ndarray."""
+        if not isinstance(value, NDArrayType):
+            return value
+
+        return np.asarray(value)

tests/patterns/numpy_type_test.py

@@ -37,12 +37,13 @@ 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"
-    )
+    data = ArrayContainer(array=np.array([1, 2, 3]))
+    af = asdf.AsdfFile({"data": data})
+    af.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'])}"
+        breakpoint()
+        assert isinstance(af.tree["data"], np.ndarray), (
+            f"Expected {type(np.ndarray)}, " f"got {type(af.tree['data'])}"
         )
-        assert np.all(af.tree["array"] == np.array([1, 2, 3]))
+        assert np.all(af.tree["data"] == np.array([1, 2, 3]))