From 615a3180567792d4a2613ff6ce0a911e37cf0500 Mon Sep 17 00:00:00 2001
From: Paul Delafosse <paul.delafosse@protonmail.com>
Date: Wed, 19 Mar 2025 09:23:10 +0100
Subject: [PATCH] feat: add enum map key

---
 schemars/src/json_schema_impls/indexmap2.rs   |  2 +-
 schemars/src/json_schema_impls/maps.rs        | 15 ++++++-
 schemars/tests/integration/main.rs            |  2 +
 .../tests/integration/map_with_enum_key.rs    | 40 +++++++++++++++++++
 ...ith_enum_key.rs~hashmap_with_enum_key.json | 23 +++++++++++
 5 files changed, 80 insertions(+), 2 deletions(-)
 create mode 100644 schemars/tests/integration/map_with_enum_key.rs
 create mode 100644 schemars/tests/integration/snapshots/schemars/tests/integration/map_with_enum_key.rs~hashmap_with_enum_key.json

diff --git a/schemars/src/json_schema_impls/indexmap2.rs b/schemars/src/json_schema_impls/indexmap2.rs
index 355864eb..a0ff6064 100644
--- a/schemars/src/json_schema_impls/indexmap2.rs
+++ b/schemars/src/json_schema_impls/indexmap2.rs
@@ -2,5 +2,5 @@ use crate::JsonSchema;
 use alloc::collections::{BTreeMap, BTreeSet};
 use indexmap2::{IndexMap, IndexSet};
 
-forward_impl!((<K, V: JsonSchema, H> JsonSchema for IndexMap<K, V, H>) => BTreeMap<K, V>);
+forward_impl!((<K: JsonSchema, V: JsonSchema, H> JsonSchema for IndexMap<K, V, H>) => BTreeMap<K, V>);
 forward_impl!((<T: JsonSchema, H> JsonSchema for IndexSet<T, H>) => BTreeSet<T>);
diff --git a/schemars/src/json_schema_impls/maps.rs b/schemars/src/json_schema_impls/maps.rs
index 86d43c8b..c24f48f5 100644
--- a/schemars/src/json_schema_impls/maps.rs
+++ b/schemars/src/json_schema_impls/maps.rs
@@ -7,6 +7,7 @@ macro_rules! map_impl {
         impl $($desc)+
         where
             V: JsonSchema,
+            K: JsonSchema,
         {
             always_inline!();
 
@@ -19,6 +20,18 @@ macro_rules! map_impl {
             }
 
             fn json_schema(generator: &mut SchemaGenerator) -> Schema {
+                let key_schema = generator.subschema_for::<K>();
+
+                if let Some(key_schema) = key_schema.as_object() {
+                    if key_schema.get("$ref").is_some() {
+                        return json_schema!({
+                            "type": "object",
+                            "additionalProperties": generator.subschema_for::<V>(),
+                            "propertyNames": generator.subschema_for::<K>(),
+                        })
+                    }
+                }
+
                 json_schema!({
                     "type": "object",
                     "additionalProperties": generator.subschema_for::<V>(),
@@ -31,4 +44,4 @@ macro_rules! map_impl {
 map_impl!(<K, V> JsonSchema for alloc::collections::BTreeMap<K, V>);
 
 #[cfg(feature = "std")]
-map_impl!(<K, V, H> JsonSchema for std::collections::HashMap<K, V, H>);
+map_impl!(<K, V, H> JsonSchema for std::collections::HashMap<K, V, H>);
\ No newline at end of file
diff --git a/schemars/tests/integration/main.rs b/schemars/tests/integration/main.rs
index eb445268..6ff25786 100644
--- a/schemars/tests/integration/main.rs
+++ b/schemars/tests/integration/main.rs
@@ -50,6 +50,8 @@ mod url;
 #[cfg(feature = "uuid1")]
 mod uuid;
 mod validator;
+#[cfg(feature = "std")]
+mod map_with_enum_key;
 
 mod prelude {
     pub(crate) use crate::test;
diff --git a/schemars/tests/integration/map_with_enum_key.rs b/schemars/tests/integration/map_with_enum_key.rs
new file mode 100644
index 00000000..9c28ce5f
--- /dev/null
+++ b/schemars/tests/integration/map_with_enum_key.rs
@@ -0,0 +1,40 @@
+use crate::prelude::*;
+use std::collections::HashMap;
+
+use schemars::JsonSchema;
+use serde::{Deserialize, Serialize};
+
+#[derive(JsonSchema, Deserialize, Serialize)]
+struct Map {
+    inner: HashMap<Key, String>,
+}
+
+#[derive(Debug, JsonSchema, Deserialize, Serialize, Eq, PartialEq, Hash, Copy, Clone)]
+#[serde(deny_unknown_fields, rename_all = "kebab-case", into = "&str")]
+enum Key {
+    A,
+    B,
+}
+
+impl From<Key> for &str {
+    fn from(value: Key) -> Self {
+        match value {
+            Key::A => "a",
+            Key::B => "b",
+        }
+    }
+}
+impl From<&str> for Key {
+    fn from(value: &str) -> Self {
+        match value {
+            "a" => Key::A,
+            "b" => Key::B,
+            _ => unreachable!(),
+        }
+    }
+}
+
+#[test]
+fn hashmap_with_enum_key() {
+    test!(Map).assert_snapshot();
+}
\ No newline at end of file
diff --git a/schemars/tests/integration/snapshots/schemars/tests/integration/map_with_enum_key.rs~hashmap_with_enum_key.json b/schemars/tests/integration/snapshots/schemars/tests/integration/map_with_enum_key.rs~hashmap_with_enum_key.json
new file mode 100644
index 00000000..a914bec1
--- /dev/null
+++ b/schemars/tests/integration/snapshots/schemars/tests/integration/map_with_enum_key.rs~hashmap_with_enum_key.json
@@ -0,0 +1,23 @@
+{
+  "$defs": {
+    "Key": {
+      "enum": ["a", "b"],
+      "type": "string"
+    }
+  },
+  "$schema": "https://json-schema.org/draft/2020-12/schema",
+  "properties": {
+    "inner": {
+      "additionalProperties": {
+        "type": "string"
+      },
+      "propertyNames": {
+        "$ref": "#/$defs/Key"
+      },
+      "type": "object"
+    }
+  },
+  "required": ["inner"],
+  "title": "Map",
+  "type": "object"
+}
\ No newline at end of file