From d81e678c65d0a7be2995926476442dbaf9cc6132 Mon Sep 17 00:00:00 2001
From: secustor <sebastian@poxhofer.at>
Date: Mon, 15 Jan 2024 22:28:59 +0100
Subject: [PATCH] chore(util/yaml): allow to provide zod schemas to YAML parser

---
 lib/util/yaml.spec.ts | 119 ++++++++++++++++++++++++++++++++++++++++++
 lib/util/yaml.ts      |  44 ++++++++++++----
 2 files changed, 153 insertions(+), 10 deletions(-)

diff --git a/lib/util/yaml.spec.ts b/lib/util/yaml.spec.ts
index b09331b7d3474b..109b5abb627936 100644
--- a/lib/util/yaml.spec.ts
+++ b/lib/util/yaml.spec.ts
@@ -1,4 +1,5 @@
 import { codeBlock } from 'common-tags';
+import { z } from 'zod';
 import { parseSingleYaml, parseYaml } from './yaml';
 
 describe('util/yaml', () => {
@@ -22,6 +23,31 @@ describe('util/yaml', () => {
       ]);
     });
 
+    it('should parse content with single document with schema', () => {
+      expect(
+        parseYaml(
+          codeBlock`
+      myObject:
+        aString: value
+      `,
+          null,
+          {
+            customSchema: z.object({
+              myObject: z.object({
+                aString: z.string(),
+              }),
+            }),
+          },
+        ),
+      ).toEqual([
+        {
+          myObject: {
+            aString: 'value',
+          },
+        },
+      ]);
+    });
+
     it('should parse content with multiple documents', () => {
       expect(
         parseYaml(codeBlock`
@@ -42,6 +68,60 @@ describe('util/yaml', () => {
       ]);
     });
 
+    it('should parse content with multiple documents with schema', () => {
+      expect(
+        parseYaml(
+          codeBlock`
+      myObject:
+        aString: foo
+      ---
+      myObject:
+        aString: bar
+      `,
+          null,
+          {
+            customSchema: z.object({
+              myObject: z.object({
+                aString: z.string(),
+              }),
+            }),
+          },
+        ),
+      ).toEqual([
+        {
+          myObject: {
+            aString: 'foo',
+          },
+        },
+        {
+          myObject: {
+            aString: 'bar',
+          },
+        },
+      ]);
+    });
+
+    it('should throw if schema does not match', () => {
+      expect(() =>
+        parseYaml(
+          codeBlock`
+      myObject:
+        aString: foo
+      ---
+      aString: bar
+      `,
+          null,
+          {
+            customSchema: z.object({
+              myObject: z.object({
+                aString: z.string(),
+              }),
+            }),
+          },
+        ),
+      ).toThrow();
+    });
+
     it('should parse content with templates', () => {
       expect(
         parseYaml(
@@ -85,6 +165,45 @@ describe('util/yaml', () => {
       });
     });
 
+    it('should parse content with single document with schema', () => {
+      expect(
+        parseSingleYaml(
+          codeBlock`
+      myObject:
+        aString: value
+      `,
+          {
+            customSchema: z.object({
+              myObject: z.object({
+                aString: z.string(),
+              }),
+            }),
+          },
+        ),
+      ).toEqual({
+        myObject: {
+          aString: 'value',
+        },
+      });
+    });
+
+    it('should throw with single document with schema if parsing fails', () => {
+      expect(() =>
+        parseSingleYaml(
+          codeBlock`
+      myObject: foo
+      `,
+          {
+            customSchema: z.object({
+              myObject: z.object({
+                aString: z.string(),
+              }),
+            }),
+          },
+        ),
+      ).toThrow();
+    });
+
     it('should parse content with multiple documents', () => {
       expect(() =>
         parseSingleYaml(codeBlock`
diff --git a/lib/util/yaml.ts b/lib/util/yaml.ts
index 8921803f2975ba..bbf146b92f0abd 100644
--- a/lib/util/yaml.ts
+++ b/lib/util/yaml.ts
@@ -5,28 +5,52 @@ import {
   load as single,
   dump as upstreamDump,
 } from 'js-yaml';
+import type { ZodType } from 'zod';
 import { regEx } from './regex';
 
-interface YamlOptions extends LoadOptions {
+type YamlOptions<
+  ResT = unknown,
+  Schema extends ZodType<ResT> = ZodType<ResT>,
+> = {
+  customSchema?: Schema;
   removeTemplates?: boolean;
-}
+} & LoadOptions;
 
-export function parseYaml(
+export function parseYaml<ResT = unknown>(
   content: string,
   iterator?: null | undefined,
-  options?: YamlOptions,
-): unknown[] {
+  options?: YamlOptions<ResT>,
+): ResT[] {
   const massagedContent = massageContent(content, options);
 
-  return multiple(massagedContent, iterator, options);
+  const rawDocuments = multiple(massagedContent, iterator, options);
+
+  const schema = options?.customSchema;
+  if (!schema) {
+    return rawDocuments as ResT[];
+  }
+
+  const parsed: ResT[] = [];
+  for (const element of rawDocuments) {
+    const singleParsed = schema.parse(element);
+    parsed.push(singleParsed);
+  }
+  return parsed;
 }
 
-export function parseSingleYaml(
+export function parseSingleYaml<ResT = unknown>(
   content: string,
-  options?: YamlOptions,
-): unknown {
+  options?: YamlOptions<ResT>,
+): ResT {
   const massagedContent = massageContent(content, options);
-  return single(massagedContent, options);
+  const rawDocument = single(massagedContent, options);
+
+  const schema = options?.customSchema;
+  if (!schema) {
+    return rawDocument as ResT;
+  }
+
+  return schema.parse(rawDocument);
 }
 
 export function dump(obj: any, opts?: DumpOptions | undefined): string {