Preserve NWB subject age-reference metadata

dandi/consts.py

@@ -32,6 +32,7 @@
 
 metadata_nwb_subject_fields = (
     "age",
+    "age__reference",
     "date_of_birth",
     "genotype",
     "sex",

dandi/metadata/util.py

@@ -247,6 +247,8 @@ def extract_age(metadata: dict) -> models.PropertyValue | None:
     except (KeyError, TypeError, ValueError):
         if metadata.get("age") is not None:
             duration, ref = parse_age(metadata["age"])
+            if metadata.get("age__reference") == "gestational":
+                ref = "Gestational"
         else:
             return None
     else:

dandi/tests/test_metadata.py

@@ -34,6 +34,7 @@
 import numpy as np
 from pydantic import ByteSize
 from pynwb import NWBHDF5IO, NWBFile, TimeSeries
+from pynwb.file import Subject
 import pytest
 import requests
 from semantic_version import Version
@@ -62,6 +63,12 @@
 
 METADATA_DIR = Path(__file__).with_name("data") / "metadata"
 
+AGE_REFERENCE_TYPES = {
+    "birth": AgeReferenceType.BirthReference,
+    "gestational": AgeReferenceType.GestationalReference,
+    None: AgeReferenceType.BirthReference,
+}
+
 mark_xfail_ontobee = pytest.mark.xfail(
     condition="not config.getoption('--scheduled')",
     reason="Flaky ontobee site",
@@ -475,6 +482,54 @@ def test_time_extract() -> None:
     )
 
 
+@pytest.mark.ai_generated
+@pytest.mark.parametrize(
+    "reference", [ref for ref in AGE_REFERENCE_TYPES.keys() if ref is not None]
+)
+def test_age_reference_roundtrip(tmp_path: Path, reference: str) -> None:
+    path = tmp_path / "age.nwb"
+    nwb = NWBFile("age reference", "age", datetime(2020, 1, 1, tzinfo=tzutc()))
+    nwb.subject = Subject(age="P3W", age__reference=reference)
+    with NWBHDF5IO(str(path), "w") as io:
+        io.write(nwb)
+    metadata = get_metadata(path)
+    assert metadata["age__reference"] == reference
+    age = extract_age(metadata)
+    assert age is not None
+    assert age.value == "P3W"
+    assert age.valueReference == PropertyValue(value=AGE_REFERENCE_TYPES[reference])
+
+
+@pytest.mark.ai_generated
+@pytest.mark.parametrize("reference", AGE_REFERENCE_TYPES.keys())
+@pytest.mark.parametrize("legacy", [False, True])
+def test_age_reference_legacy(reference: str | None, legacy: bool) -> None:
+    age = extract_age(
+        {"age": "Gestational week 3" if legacy else "P3W", "age__reference": reference}
+    )
+    assert age is not None
+    assert age.value == "P3W"
+    assert age.valueReference == PropertyValue(
+        value=AGE_REFERENCE_TYPES["gestational" if legacy else reference]
+    )
+
+
+@pytest.mark.ai_generated
+@pytest.mark.parametrize("reference", AGE_REFERENCE_TYPES.keys())
+def test_age_reference_preserves_birth_dates(reference: str | None) -> None:
+    age = extract_age(
+        {
+            "age": "P3W",
+            "age__reference": reference,
+            "date_of_birth": "2020-01-01T00:00:00Z",
+            "session_start_time": "2020-01-03T00:00:00Z",
+        }
+    )
+    assert age is not None
+    assert age.value == "P2D"
+    assert age.valueReference == PropertyValue(value=AGE_REFERENCE_TYPES["birth"])
+
+
 def test_time_extract_gest() -> None:
     """extract age with Gestational ref"""
     meta_birth = {