From e1f6c90fc0733b2eec9f400e2359fadc497904c1 Mon Sep 17 00:00:00 2001
From: Elio Struyf <elio@struyfconsulting.be>
Date: Sat, 11 Nov 2023 10:38:31 -0800
Subject: [PATCH] Enhancement: Better content type creation for Astro #704

---
 ssg-scripts/astro.collections.mjs | 50 +++++++++++++++++++++++++++++--
 1 file changed, 48 insertions(+), 2 deletions(-)

diff --git a/ssg-scripts/astro.collections.mjs b/ssg-scripts/astro.collections.mjs
index 947dd9dc..b96d59a7 100644
--- a/ssg-scripts/astro.collections.mjs
+++ b/ssg-scripts/astro.collections.mjs
@@ -3,7 +3,8 @@ import { join } from 'path';
 import { createServer } from 'vite';
 import zod from 'astro/zod';
 
-const { ZodDefault, ZodObject, ZodOptional, ZodString, ZodEffects, ZodEnum } = zod;
+const { ZodDefault, ZodObject, ZodOptional, ZodString, ZodEffects, ZodEnum, ZodUnion, ZodArray } =
+  zod;
 
 /**
  * Process the Zod field
@@ -21,6 +22,17 @@ function getField(field, isOptional = false, defaultValue = undefined) {
     return getField(type, isOptional, defaultValue);
   }
 
+  if (field instanceof ZodEffects) {
+    const type = field.sourceType();
+    return getField(type, isOptional, defaultValue);
+  }
+
+  if (field instanceof ZodUnion) {
+    const types = field._def.options;
+    const type = types[0];
+    return getField(type, isOptional, defaultValue);
+  }
+
   if (field instanceof ZodDefault) {
     // https://github.com/colinhacks/zod/blob/master/README.md#default
     isOptional = true;
@@ -61,6 +73,26 @@ function generateFieldInfo(name, type) {
     required: !isFieldOptional
   };
 
+  switch (fieldInfo.type) {
+    case 'ZodString':
+      fieldInfo.type = 'string';
+      break;
+    case 'ZodNumber':
+      fieldInfo.type = 'number';
+      break;
+    case 'ZodBoolean':
+      fieldInfo.type = 'boolean';
+      break;
+    case 'ZodDate':
+      fieldInfo.type = 'datetime';
+      break;
+    case 'ZodArray':
+      fieldInfo.type = 'choice';
+      break;
+    default:
+      break;
+  }
+
   if (fieldType instanceof ZodObject) {
     const subFields = extractFieldInfoFromShape(fieldType);
 
@@ -71,10 +103,24 @@ function generateFieldInfo(name, type) {
     fieldInfo.type = fieldType.sourceType().fmFieldType;
   }
 
+  if (fieldInfo.name.toLowerCase().includes('image')) {
+    fieldInfo.type = 'image';
+  }
+
   if (fieldType instanceof ZodEnum) {
     fieldInfo.options = fieldType.options;
   }
 
+  if (fieldType instanceof ZodArray) {
+    if (fieldInfo.name === 'tags') {
+      fieldInfo.type = 'tags';
+    } else if (fieldInfo.name === 'categories') {
+      fieldInfo.type = 'categories';
+    } else {
+      fieldInfo.options = fieldType.options;
+    }
+  }
+
   if (fieldType instanceof ZodString) {
     // https://github.com/StefanTerdell/zod-to-json-schema/blob/master/src/parsers/string.ts#L45
     if (fieldType._def.checks && fieldType._def.checks.length > 0) {
@@ -205,7 +251,7 @@ const astroContentModulePlugin = () => {
 
     writeFileSync(
       join(process.cwd(), `./.frontmatter/temp/astro.collections.json`),
-      JSON.stringify(processedCollections)
+      JSON.stringify(processedCollections, null, 2)
     );
 
     console.log('Collections generated successfully');