From f4750de0da3113d917feffeaaf9f26655b51104c Mon Sep 17 00:00:00 2001
From: Charles Korn <me@charleskorn.com>
Date: Fri, 23 Oct 2020 11:50:57 +1100
Subject: [PATCH] Add method to get the key node for an entry in a YamlMap.

---
 .../kotlin/com/charleskorn/kaml/YamlNode.kt   |  2 ++
 .../com/charleskorn/kaml/YamlMapTest.kt       | 27 +++++++++++++++++++
 2 files changed, 29 insertions(+)

diff --git a/src/main/kotlin/com/charleskorn/kaml/YamlNode.kt b/src/main/kotlin/com/charleskorn/kaml/YamlNode.kt
index ae8b64db..f1114e1b 100644
--- a/src/main/kotlin/com/charleskorn/kaml/YamlNode.kt
+++ b/src/main/kotlin/com/charleskorn/kaml/YamlNode.kt
@@ -192,6 +192,8 @@ public data class YamlMap(val entries: Map<YamlScalar, YamlNode>, override val p
         else -> throw IncorrectTypeException("Value for '$key' is not a scalar.", node.path)
     }
 
+    public fun getKey(key: String): YamlScalar? = entries.keys.singleOrNull { it.content == key }
+
     override fun withPath(newPath: YamlPath): YamlMap {
         val updatedEntries = entries
             .mapKeys { (k, _) -> k.withPath(replacePathOnChild(k, newPath)) }
diff --git a/src/test/kotlin/com/charleskorn/kaml/YamlMapTest.kt b/src/test/kotlin/com/charleskorn/kaml/YamlMapTest.kt
index d90170fd..fb910188 100644
--- a/src/test/kotlin/com/charleskorn/kaml/YamlMapTest.kt
+++ b/src/test/kotlin/com/charleskorn/kaml/YamlMapTest.kt
@@ -299,6 +299,33 @@ object YamlMapTest : Spek({
             }
         }
 
+        describe("getting keys of the map") {
+            val helloKeyPath = YamlPath.root.withMapElementKey("hello", Location(1, 1))
+            val helloValuePath = helloKeyPath.withMapElementValue(Location(2, 1))
+            val alsoKeyPath = YamlPath.root.withMapElementKey("also", Location(3, 1))
+            val alsoValuePath = alsoKeyPath.withMapElementValue(Location(4, 1))
+
+            val map = YamlMap(
+                mapOf(
+                    YamlScalar("hello", helloKeyPath) to YamlScalar("world", helloValuePath),
+                    YamlScalar("also", alsoKeyPath) to YamlScalar("something", alsoValuePath)
+                ),
+                YamlPath.root
+            )
+
+            context("the key is not in the map") {
+                it("returns null") {
+                    expect(map.getKey("something else")).toBe(null)
+                }
+            }
+
+            context("the key is in the map") {
+                it("returns the node for that key") {
+                    expect(map.getKey("hello")).toBe(YamlScalar("hello", helloKeyPath))
+                }
+            }
+        }
+
         describe("replacing its path") {
             val originalPath = YamlPath.root
             val originalKey1Path = originalPath.withMapElementKey("key1", Location(4, 1))