From 21962f90395fde84bee43e31d0a69fe63e066313 Mon Sep 17 00:00:00 2001
From: "Keto D. Zhang" <keto.zhang@gmail.com>
Date: Sun, 25 Aug 2024 18:33:11 -0700
Subject: [PATCH] fix: ordering of generated schema

---
 .pre-commit-config.yaml |  6 ++++++
 asdf_pydantic/model.py  |  2 +-
 asdf_pydantic/schema.py | 22 +++++++++++++++++++++-
 tests/test_model.py     | 17 ++++++++++++++++-
 4 files changed, 44 insertions(+), 3 deletions(-)

diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 29196e9..2e401e3 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -21,6 +21,12 @@ repos:
     hooks:
       - id: black
 
+  - repo: https://github.com/compilerla/conventional-pre-commit
+    rev: v3.4.0
+    hooks:
+      - id: conventional-pre-commit
+        stages: [commit-msg]
+
 ci:
     # autofix_commit_msg: |
     #     [pre-commit.ci] auto fixes from pre-commit.com hooks
diff --git a/asdf_pydantic/model.py b/asdf_pydantic/model.py
index 09663a0..b7c3f42 100644
--- a/asdf_pydantic/model.py
+++ b/asdf_pydantic/model.py
@@ -72,7 +72,7 @@ def model_asdf_schema(
         )
         json_schema = schema_generator_instance.generate(cls.__pydantic_core_schema__)
 
-        return f"%YAML 1.1\n---\n{yaml.safe_dump(json_schema)}"
+        return f"%YAML 1.1\n---\n{yaml.safe_dump(json_schema, sort_keys=False)}"
 
     @classmethod
     @deprecated(
diff --git a/asdf_pydantic/schema.py b/asdf_pydantic/schema.py
index 952a695..2adc6ab 100644
--- a/asdf_pydantic/schema.py
+++ b/asdf_pydantic/schema.py
@@ -3,6 +3,17 @@
 from pydantic.json_schema import GenerateJsonSchema
 
 DEFAULT_ASDF_SCHEMA_REF_TEMPLATE = "#/definitions/{model}"
+DESIRED_ASDF_SCHEMA_KEY_ORDER = (
+    "$schema",
+    "id",
+    "title",
+    "type",
+    "properties",
+    "allOf",
+    "anyOf",
+    "required",
+    "$defs",
+)
 
 
 class GenerateAsdfSchema(GenerateJsonSchema):
@@ -33,7 +44,16 @@ def generate(self, schema, mode="validation"):
         if self.tag_uri:
             json_schema["$schema"] = self.schema_dialect
             json_schema["id"] = f"{self.tag_uri}/schema"
-            json_schema["tag"] = f"tag:{self.tag_uri.split('://', maxsplit=2)[-1]}"
+
+        # Order keys
+        json_schema = {
+            **{
+                key: json_schema[key]
+                for key in DESIRED_ASDF_SCHEMA_KEY_ORDER
+                if key in json_schema
+            },
+            **json_schema,  # Rest of the keys not in order list
+        }
 
         # TODO: Convert jsonschema 2020-12 to ASDF schema
         return json_schema
diff --git a/tests/test_model.py b/tests/test_model.py
index 29bb5fe..9f78227 100644
--- a/tests/test_model.py
+++ b/tests/test_model.py
@@ -52,8 +52,23 @@ class TestModel(AsdfPydanticModel):
 
 
 ########################################################################################
-# SCHEMA
+# GENERATED SCHEMA
 ########################################################################################
+def test_generated_schema_keys_in_order():
+    class TestModel(AsdfPydanticModel):
+        _tag = "asdf://asdf-pydantic/tags/test-0.0.1"
+        foo: str
+
+    assert list(yaml.safe_load(TestModel.model_asdf_schema()).keys()) == [
+        "$schema",
+        "id",
+        "title",
+        "type",
+        "properties",
+        "required",
+    ]
+
+
 def test_generated_schema_id_uses_tag_in_pattern():
     class TestModel(AsdfPydanticModel):
         _tag = "asdf://asdf-pydantic/tags/test-0.0.1"