[ENH] annotation order (#1220)

* wip: add order to note_keys

* add utils and fix style

Author: jdkent

store/backend/neurostore/ingest/__init__.py

@@ -364,7 +364,8 @@ def ingest_neurosynth(max_rows=None):
 
         # add notes to annotation
         annot.note_keys = {
-            k: _check_type(v) for k, v in annotation_row._asdict().items()
+            k: {"type": _check_type(v) or "string", "order": idx}
+            for idx, (k, v) in enumerate(annotation_row._asdict().items())
         }
         annot.annotation_analyses = notes
         for note in notes:

store/backend/neurostore/resources/data.py

@@ -393,6 +393,84 @@ def get_affected_ids(self, ids):
         }
         return unique_ids
 
+    @staticmethod
+    def _ordered_note_keys(note_keys):
+        if not note_keys:
+            return []
+        keys = list(note_keys.keys())
+        alphabetical = sorted(keys)
+        if keys == alphabetical:
+            return alphabetical
+        return keys
+
+    @classmethod
+    def _normalize_note_keys(cls, note_keys):
+        if note_keys is None:
+            return None
+        if not isinstance(note_keys, dict):
+            abort_validation("`note_keys` must be an object.")
+
+        ordered_keys = cls._ordered_note_keys(note_keys)
+        normalized = OrderedDict()
+        used_orders = set()
+        next_order = 0
+
+        for key in ordered_keys:
+            descriptor = note_keys.get(key) or {}
+            note_type = descriptor.get("type")
+            if note_type not in {"string", "number", "boolean"}:
+                abort_validation(
+                    "Invalid `type` for note_keys entry "
+                    f"'{key}', choose from: ['boolean', 'number', 'string']."
+                )
+
+            order = descriptor.get("order")
+            if isinstance(order, bool) or (
+                order is not None and not isinstance(order, int)
+            ):
+                order = None
+
+            if isinstance(order, int) and order not in used_orders:
+                used_orders.add(order)
+                if order >= next_order:
+                    next_order = order + 1
+            else:
+                while next_order in used_orders:
+                    next_order += 1
+                order = next_order
+                used_orders.add(order)
+                next_order += 1
+
+            normalized[key] = {"type": note_type, "order": order}
+
+        return normalized
+
+    @classmethod
+    def _merge_note_keys(cls, existing, additions):
+        """
+        additions is a mapping of key -> type
+        """
+        base = cls._normalize_note_keys(existing or {}) or OrderedDict()
+        used_orders = {v.get("order") for v in base.values() if isinstance(v, dict)}
+        used_orders = {o for o in used_orders if isinstance(o, int)}
+        next_order = max(used_orders, default=-1) + 1
+
+        for key, value_type in additions.items():
+            if key in base:
+                descriptor = base[key] or {}
+                descriptor["type"] = value_type or descriptor.get("type") or "string"
+                base[key] = descriptor
+                continue
+
+            descriptor = {
+                "type": value_type or "string",
+                "order": next_order,
+            }
+            base[key] = descriptor
+            next_order += 1
+
+        return base
+
     @classmethod
     def load_nested_records(cls, data, record=None):
         if not data:
@@ -554,6 +632,9 @@ def put(self, id):
         schema = self._schema()
         data = schema.load(request_data)
 
+        if "note_keys" in data:
+            data["note_keys"] = self._normalize_note_keys(data["note_keys"])
+
         pipeline_payload = data.pop("pipelines", [])
 
         args = {}
@@ -942,12 +1023,10 @@ def _apply_pipeline_columns(self, annotation, data, specs, column_counter):
 
         if column_types:
             if data.get("note_keys") is None:
-                note_keys = dict(annotation.note_keys or {})
+                note_keys = self._normalize_note_keys(annotation.note_keys or {})
             else:
-                note_keys = dict(data["note_keys"])
-            for key, value_type in column_types.items():
-                note_keys[key] = value_type or "string"
-            data["note_keys"] = note_keys
+                note_keys = self._normalize_note_keys(data["note_keys"])
+            data["note_keys"] = self._merge_note_keys(note_keys, column_types)
 
         data["annotation_analyses"] = list(note_map.values())
 

store/backend/neurostore/schemas/data.py

@@ -663,6 +663,70 @@ class AnnotationPipelineSchema(BaseSchema):
     columns = fields.List(fields.String(), required=True)
 
 
+class NoteKeysField(fields.Field):
+    allowed_types = {"string", "number", "boolean"}
+
+    def _serialize(self, value, attr, obj, **kwargs):
+        if not value:
+            return {}
+        serialized = {}
+        for key, descriptor in value.items():
+            if not isinstance(descriptor, dict):
+                continue
+            serialized[key] = {
+                "type": descriptor.get("type"),
+                "order": descriptor.get("order"),
+            }
+        return serialized
+
+    def _deserialize(self, value, attr, data, **kwargs):
+        if value is None:
+            return {}
+        if not isinstance(value, dict):
+            raise ValidationError("`note_keys` must be an object.")
+
+        normalized = {}
+        used_orders = set()
+        explicit_orders = []
+        for descriptor in value.values():
+            if isinstance(descriptor, dict) and isinstance(
+                descriptor.get("order"), int
+            ):
+                explicit_orders.append(descriptor["order"])
+        next_order = max(explicit_orders, default=-1) + 1
+
+        for key, descriptor in value.items():
+            if not isinstance(descriptor, dict):
+                raise ValidationError("Each note key must map to an object.")
+
+            note_type = descriptor.get("type")
+            if note_type not in self.allowed_types:
+                raise ValidationError(
+                    f"Invalid note type for '{key}', choose from: {sorted(self.allowed_types)}"
+                )
+
+            order = descriptor.get("order")
+            if isinstance(order, bool) or (
+                order is not None and not isinstance(order, int)
+            ):
+                order = None
+
+            if isinstance(order, int) and order not in used_orders:
+                used_orders.add(order)
+                if order >= next_order:
+                    next_order = order + 1
+            else:
+                while next_order in used_orders:
+                    next_order += 1
+                order = next_order
+                used_orders.add(order)
+                next_order += 1
+
+            normalized[key] = {"type": note_type, "order": order}
+
+        return normalized
+
+
 class AnnotationSchema(BaseDataSchema):
     # serialization
     studyset_id = fields.String(data_key="studyset")
@@ -675,7 +739,7 @@ class AnnotationSchema(BaseDataSchema):
     source_id = fields.String(dump_only=True, allow_none=True)
     source_updated_at = fields.DateTime(dump_only=True, allow_none=True)
 
-    note_keys = fields.Dict()
+    note_keys = NoteKeysField()
     metadata = fields.Dict(attribute="metadata_", dump_only=True)
     # deserialization
     metadata_ = fields.Dict(data_key="metadata", load_only=True, allow_none=True)

store/backend/neurostore/tests/api/test_annotations.py

@@ -13,6 +13,7 @@
     PipelineConfig,
     PipelineStudyResult,
 )
+from ..utils import ordered_note_keys
 
 
 def _create_annotation_with_two_analyses(session, user):
@@ -39,7 +40,7 @@ def _create_annotation_with_two_analyses(session, user):
         name="Test Annotation",
         studyset=studyset,
         user=user,
-        note_keys={"existing": "string"},
+        note_keys=ordered_note_keys({"existing": "string"}),
     )
 
     session.add(annotation)
@@ -92,7 +93,7 @@ def test_blank_annotation_populates_note_fields(
     auth_client, ingest_neurosynth, session
 ):
     dset = Studyset.query.first()
-    note_keys = {"included": "boolean", "quality": "string"}
+    note_keys = ordered_note_keys({"included": "boolean", "quality": "string"})
     payload = {
         "studyset": dset.id,
         "note_keys": note_keys,
@@ -121,7 +122,7 @@ def test_annotation_rejects_empty_note(auth_client, ingest_neurosynth, session):
                 "note": {},
             }
         ],
-        "note_keys": {"included": "boolean"},
+        "note_keys": ordered_note_keys({"included": "boolean"}),
         "name": "invalid annotation",
     }
 
@@ -143,7 +144,7 @@ def test_post_annotation(auth_client, ingest_neurosynth, session):
     payload = {
         "studyset": dset.id,
         "notes": data,
-        "note_keys": {"foo": "string"},
+        "note_keys": ordered_note_keys({"foo": "string"}),
         "name": "mah notes",
     }
     resp = auth_client.post("/api/annotations/", data=payload)
@@ -280,7 +281,7 @@ def test_blank_slate_creation(auth_client, session):
     # create annotation
     annotation_data = {
         "studyset": ss_id,
-        "note_keys": {"include": "boolean"},
+        "note_keys": ordered_note_keys({"include": "boolean"}),
         "name": "mah notes",
     }
     annotation_post = auth_client.post("/api/annotations/", data=annotation_data)
@@ -342,7 +343,7 @@ def test_mismatched_notes(auth_client, ingest_neurosynth, session):
         for s in dset.studies
         for a in s.analyses
     ]
-    note_keys = {"foo": "string", "doo": "string"}
+    note_keys = ordered_note_keys({"foo": "string", "doo": "string"})
     payload = {
         "studyset": dset.id,
         "notes": data,
@@ -377,7 +378,7 @@ def test_put_nonexistent_analysis(auth_client, ingest_neurosynth, session):
         for s in dset.studies
         for a in s.analyses
     ]
-    note_keys = {"foo": "string", "doo": "string"}
+    note_keys = ordered_note_keys({"foo": "string", "doo": "string"})
     payload = {
         "studyset": dset.id,
         "notes": data,
@@ -412,7 +413,7 @@ def test_post_put_subset_of_analyses(auth_client, ingest_neurosynth, session):
     # remove last note
     data.pop()
 
-    note_keys = {"foo": "string", "doo": "string"}
+    note_keys = ordered_note_keys({"foo": "string", "doo": "string"})
     payload = {
         "studyset": dset.id,
         "notes": data,
@@ -444,7 +445,7 @@ def test_correct_note_overwrite(auth_client, ingest_neurosynth, session):
     payload = {
         "studyset": dset.id,
         "notes": data,
-        "note_keys": {"foo": "string", "doo": "string"},
+        "note_keys": ordered_note_keys({"foo": "string", "doo": "string"}),
         "name": "mah notes",
     }
 
@@ -514,10 +515,10 @@ def test_put_annotation_applies_pipeline_columns(auth_client, session):
     assert resp.status_code == 200
     body = resp.json()
 
-    assert body["note_keys"]["existing"] == "string"
-    assert body["note_keys"]["string_field"] == "string"
-    assert body["note_keys"]["numeric_field"] == "number"
-    assert body["note_keys"]["name"] == "string"
+    assert body["note_keys"]["existing"]["type"] == "string"
+    assert body["note_keys"]["string_field"]["type"] == "string"
+    assert body["note_keys"]["numeric_field"]["type"] == "number"
+    assert body["note_keys"]["name"]["type"] == "string"
 
     notes = body["notes"]
     assert len(notes) == 2
@@ -603,9 +604,9 @@ def test_put_annotation_pipeline_column_conflict_suffix(auth_client, session):
 
     assert key_one in body["note_keys"]
     assert key_two in body["note_keys"]
-    assert body["note_keys"][key_one] == "string"
-    assert body["note_keys"][key_two] == "string"
-    assert body["note_keys"]["name"] == "string"
+    assert body["note_keys"][key_one]["type"] == "string"
+    assert body["note_keys"][key_two]["type"] == "string"
+    assert body["note_keys"]["name"]["type"] == "string"
 
     for entry in body["notes"]:
         note = entry["note"]
@@ -625,7 +626,7 @@ def test_annotation_analyses_post(auth_client, ingest_neurosynth, session):
     payload = {
         "studyset": dset.id,
         "notes": data,
-        "note_keys": {"foo": "string", "doo": "string"},
+        "note_keys": ordered_note_keys({"foo": "string", "doo": "string"}),
         "name": "mah notes",
     }
 

store/backend/neurostore/tests/api/test_studysets.py

@@ -164,7 +164,7 @@ def _create_studyset_with_annotation(auth_client, study_ids, name="clone-source"
 
     annotation_payload = {
         "studyset": studyset_id,
-        "note_keys": {"include": "boolean"},
+        "note_keys": {"include": {"type": "boolean", "order": 0}},
         "name": "annotation for clone",
     }
     annotation_resp = auth_client.post("/api/annotations/", data=annotation_payload)

store/backend/neurostore/tests/conftest.py

@@ -29,6 +29,7 @@
 import vcr
 
 import logging
+from .utils import ordered_note_keys
 
 LOGGER = logging.getLogger(__name__)
 
@@ -602,7 +603,7 @@ def user_data(session, mock_add_users):
             annotation = Annotation(
                 name=name + "annotation",
                 source="neurostore",
-                note_keys={"food": "string"},
+                note_keys=ordered_note_keys({"food": "string"}),
                 studyset=studyset,
                 user=user,
             )
@@ -639,13 +640,13 @@ def simple_neurosynth_annotation(session, ingest_neurosynth):
                 )
             )
 
-        smol_annot = Annotation(
-            name="smol " + annot.name,
-            source="neurostore",
-            studyset=annot.studyset,
-            note_keys={"animal": "number"},
-            annotation_analyses=smol_notes,
-        )
+            smol_annot = Annotation(
+                name="smol " + annot.name,
+                source="neurostore",
+                studyset=annot.studyset,
+                note_keys=ordered_note_keys({"animal": "number"}),
+                annotation_analyses=smol_notes,
+            )
     session.add(smol_annot)
     session.commit()
 

store/backend/neurostore/tests/utils.py

@@ -0,0 +1,8 @@
+def ordered_note_keys(type_map):
+    """
+    Build ordered note key descriptors from a simple name->type mapping.
+    """
+    return {
+        key: {"type": value_type, "order": idx}
+        for idx, (key, value_type) in enumerate(type_map.items())
+    }