Add `extend` options to docs and i18n schemas (#1162)

6 files changed, 289 insertions, 98 deletions

diff --git a/.changeset/short-toes-cheat.md b/.changeset/short-toes-cheat.md
new file mode 100644
index 00000000..b62a37a8
--- /dev/null
+++ b/.changeset/short-toes-cheat.md
@@ -0,0 +1,5 @@
+---
+'@astrojs/starlight': minor
+---
+
+Adds support for extending Starlight’s content collection schemas
diff --git a/docs/src/content/docs/guides/authoring-content.md b/docs/src/content/docs/guides/authoring-content.md
index 10a4e409..f377fb07 100644
--- a/docs/src/content/docs/guides/authoring-content.md
+++ b/docs/src/content/docs/guides/authoring-content.md
@@ -7,6 +7,22 @@ Starlight supports the full range of [Markdown](https://daringfireball.net/proje
 
 Please be sure to check the [MDX docs](https://mdxjs.com/docs/what-is-mdx/#markdown) or [Markdoc docs](https://markdoc.dev/docs/syntax) if using those file formats, as Markdown support and usage can differ.
 
+## Frontmatter
+
+You can customize individual pages in Starlight by setting values in their frontmatter.
+Frontmatter is set at the top of your files between `---` separators:
+
+```md title="src/content/docs/example.md"
+---
+title: My page title
+---
+
+Page content follows the second `---`.
+```
+
+Every page must include at least a `title`.
+See the [frontmatter reference](/reference/frontmatter/) for all available fields and how to add custom fields.
+
 ## Inline styles
 
 Text can be **bold**, _italic_, or ~~strikethrough~~.
diff --git a/docs/src/content/docs/guides/i18n.mdx b/docs/src/content/docs/guides/i18n.mdx
index ac5b4806..5b573ac4 100644
--- a/docs/src/content/docs/guides/i18n.mdx
+++ b/docs/src/content/docs/guides/i18n.mdx
@@ -233,3 +233,28 @@ You can provide translations for additional languages you support — or overrid
    	"pagefind.searching": "Searching for [SEARCH_TERM]..."
    }
    ```
+
+### Extend translation schema
+
+Add custom keys to your site’s translation dictionaries by setting `extend` in the `i18nSchema()` options.
+In the following example, a new, optional `custom.label` key is added to the default keys:
+
+```diff lang="js"
+// src/content/config.ts
+import { defineCollection, z } from 'astro:content';
+import { docsSchema, i18nSchema } from '@astrojs/starlight/schema';
+
+export const collections = {
+	docs: defineCollection({ schema: docsSchema() }),
+	i18n: defineCollection({
+		type: 'data',
+		schema: i18nSchema({
++			extend: z.object({
++				'custom.label': z.string().optional(),
++			}),
+		}),
+	}),
+};
+```
+
+Learn more about content collection schemas in [“Defining a collection schema”](https://docs.astro.build/en/guides/content-collections/#defining-a-collection-schema) in the Astro docs.
diff --git a/docs/src/content/docs/reference/frontmatter.md b/docs/src/content/docs/reference/frontmatter.md
index 487b6a29..466ee614 100644
--- a/docs/src/content/docs/reference/frontmatter.md
+++ b/docs/src/content/docs/reference/frontmatter.md
@@ -370,3 +370,71 @@ sidebar:
     target: _blank
 ---
 ```
+
+## Customize frontmatter schema
+
+The frontmatter schema for Starlight’s `docs` content collection is configured in `src/content/config.ts` using the `docsSchema()` helper:
+
+```ts {3,6}
+// src/content/config.ts
+import { defineCollection } from 'astro:content';
+import { docsSchema } from '@astrojs/starlight/schema';
+
+export const collections = {
+  docs: defineCollection({ schema: docsSchema() }),
+};
+```
+
+Learn more about content collection schemas in [“Defining a collection schema”](https://docs.astro.build/en/guides/content-collections/#defining-a-collection-schema) in the Astro docs.
+
+`docsSchema()` takes the following options:
+
+### `extend`
+
+**type:** Zod schema or function that returns a Zod schema  
+**default:** `z.object({})`
+
+Extend Starlight’s schema with additional fields by setting `extend` in the `docsSchema()` options.
+The value should be a [Zod schema](https://docs.astro.build/en/guides/content-collections/#defining-datatypes-with-zod).
+
+In the following example, we provide a stricter type for `description` to make it required and add a new optional `category` field:
+
+```ts {8-13}
+// src/content/config.ts
+import { defineCollection, z } from 'astro:content';
+import { docsSchema } from '@astrojs/starlight/schema';
+
+export const collections = {
+  docs: defineCollection({
+    schema: docsSchema({
+      extend: z.object({
+        // Make a built-in field required instead of optional.
+        description: z.string(),
+        // Add a new field to the schema.
+        category: z.enum(['tutorial', 'guide', 'reference']).optional(),
+      }),
+    }),
+  }),
+};
+```
+
+To take advantage of the [Astro `image()` helper](https://docs.astro.build/en/guides/images/#images-in-content-collections), use a function that returns your schema extension:
+
+```ts {8-13}
+// src/content/config.ts
+import { defineCollection, z } from 'astro:content';
+import { docsSchema } from '@astrojs/starlight/schema';
+
+export const collections = {
+  docs: defineCollection({
+    schema: docsSchema({
+      extend: ({ image }) => {
+        return z.object({
+          // Add a field that must resolve to a local image.
+          cover: image(),
+        });
+      },
+    }),
+  }),
+};
+```
diff --git a/packages/starlight/schema.ts b/packages/starlight/schema.ts
index cece5949..a3534af0 100644
--- a/packages/starlight/schema.ts
+++ b/packages/starlight/schema.ts
@@ -8,100 +8,153 @@ import { HeroSchema } from './schemas/hero';
 import { SidebarLinkItemHTMLAttributesSchema } from './schemas/sidebar';
 export { i18nSchema } from './schemas/i18n';
 
-export function docsSchema() {
-	return (context: SchemaContext) =>
-		z.object({
-			/** The title of the current page. Required. */
-			title: z.string(),
-
-			/**
-			 * A short description of the current page’s content. Optional, but recommended.
-			 * A good description is 150–160 characters long and outlines the key content
-			 * of the page in a clear and engaging way.
-			 */
-			description: z.string().optional(),
-
-			/**
-			 * Custom URL where a reader can edit this page.
-			 * Overrides the `editLink.baseUrl` global config if set.
-			 *
-			 * Can also be set to `false` to disable showing an edit link on this page.
-			 */
-			editUrl: z.union([z.string().url(), z.boolean()]).optional().default(true),
-
-			/** Set custom `<head>` tags just for this page. */
-			head: HeadConfigSchema(),
-
-			/** Override global table of contents configuration for this page. */
-			tableOfContents: TableOfContentsSchema().optional(),
-
-			/**
-			 * Set the layout style for this page.
-			 * Can be `'doc'` (the default) or `'splash'` for a wider layout without any sidebars.
-			 */
-			template: z.enum(['doc', 'splash']).default('doc'),
-
-			/** Display a hero section on this page. */
-			hero: HeroSchema(context).optional(),
-
-			/**
-			 * The last update date of the current page.
-			 * Overrides the `lastUpdated` global config or the date generated from the Git history.
-			 */
-			lastUpdated: z.union([z.date(), z.boolean()]).optional(),
-
-			/**
-			 * The previous navigation link configuration.
-			 * Overrides the `pagination` global config or the link text and/or URL.
-			 */
-			prev: PrevNextLinkConfigSchema(),
-			/**
-			 * The next navigation link configuration.
-			 * Overrides the `pagination` global config or the link text and/or URL.
-			 */
-			next: PrevNextLinkConfigSchema(),
-
-			sidebar: z
-				.object({
-					/**
-					 * The order of this page in the navigation.
-					 * Pages are sorted by this value in ascending order. Then by slug.
-					 * If not provided, pages will be sorted alphabetically by slug.
-					 * If two pages have the same order value, they will be sorted alphabetically by slug.
-					 */
-					order: z.number().optional(),
-
-					/**
-					 * The label for this page in the navigation.
-					 * Defaults to the page `title` if not set.
-					 */
-					label: z.string().optional(),
-
-					/**
-					 * Prevents this page from being included in autogenerated sidebar groups.
-					 */
-					hidden: z.boolean().default(false),
-					/**
-					 * Adds a badge to the sidebar link.
-					 * Can be a string or an object with a variant and text.
-					 * Variants include 'note', 'tip', 'caution', 'danger', 'success', and 'default'.
-					 * Passing only a string defaults to the 'default' variant which uses the site accent color.
-					 */
-					badge: BadgeConfigSchema(),
-					/** HTML attributes to add to the sidebar link. */
-					attrs: SidebarLinkItemHTMLAttributesSchema(),
-				})
-				.default({}),
-
-			/** Display an announcement banner at the top of this page. */
-			banner: z
-				.object({
-					/** The content of the banner. Supports HTML syntax. */
-					content: z.string(),
-				})
-				.optional(),
-
-			/** Pagefind indexing for this page - set to false to disable. */
-			pagefind: z.boolean().default(true),
-		});
+/** Default content collection schema for Starlight’s `docs` collection. */
+const StarlightFrontmatterSchema = (context: SchemaContext) =>
+	z.object({
+		/** The title of the current page. Required. */
+		title: z.string(),
+
+		/**
+		 * A short description of the current page’s content. Optional, but recommended.
+		 * A good description is 150–160 characters long and outlines the key content
+		 * of the page in a clear and engaging way.
+		 */
+		description: z.string().optional(),
+
+		/**
+		 * Custom URL where a reader can edit this page.
+		 * Overrides the `editLink.baseUrl` global config if set.
+		 *
+		 * Can also be set to `false` to disable showing an edit link on this page.
+		 */
+		editUrl: z.union([z.string().url(), z.boolean()]).optional().default(true),
+
+		/** Set custom `<head>` tags just for this page. */
+		head: HeadConfigSchema(),
+
+		/** Override global table of contents configuration for this page. */
+		tableOfContents: TableOfContentsSchema().optional(),
+
+		/**
+		 * Set the layout style for this page.
+		 * Can be `'doc'` (the default) or `'splash'` for a wider layout without any sidebars.
+		 */
+		template: z.enum(['doc', 'splash']).default('doc'),
+
+		/** Display a hero section on this page. */
+		hero: HeroSchema(context).optional(),
+
+		/**
+		 * The last update date of the current page.
+		 * Overrides the `lastUpdated` global config or the date generated from the Git history.
+		 */
+		lastUpdated: z.union([z.date(), z.boolean()]).optional(),
+
+		/**
+		 * The previous navigation link configuration.
+		 * Overrides the `pagination` global config or the link text and/or URL.
+		 */
+		prev: PrevNextLinkConfigSchema(),
+		/**
+		 * The next navigation link configuration.
+		 * Overrides the `pagination` global config or the link text and/or URL.
+		 */
+		next: PrevNextLinkConfigSchema(),
+
+		sidebar: z
+			.object({
+				/**
+				 * The order of this page in the navigation.
+				 * Pages are sorted by this value in ascending order. Then by slug.
+				 * If not provided, pages will be sorted alphabetically by slug.
+				 * If two pages have the same order value, they will be sorted alphabetically by slug.
+				 */
+				order: z.number().optional(),
+
+				/**
+				 * The label for this page in the navigation.
+				 * Defaults to the page `title` if not set.
+				 */
+				label: z.string().optional(),
+
+				/**
+				 * Prevents this page from being included in autogenerated sidebar groups.
+				 */
+				hidden: z.boolean().default(false),
+				/**
+				 * Adds a badge to the sidebar link.
+				 * Can be a string or an object with a variant and text.
+				 * Variants include 'note', 'tip', 'caution', 'danger', 'success', and 'default'.
+				 * Passing only a string defaults to the 'default' variant which uses the site accent color.
+				 */
+				badge: BadgeConfigSchema(),
+				/** HTML attributes to add to the sidebar link. */
+				attrs: SidebarLinkItemHTMLAttributesSchema(),
+			})
+			.default({}),
+
+		/** Display an announcement banner at the top of this page. */
+		banner: z
+			.object({
+				/** The content of the banner. Supports HTML syntax. */
+				content: z.string(),
+			})
+			.optional(),
+
+		/** Pagefind indexing for this page - set to false to disable. */
+		pagefind: z.boolean().default(true),
+	});
+/** Type of Starlight’s default frontmatter schema. */
+type DefaultSchema = ReturnType<typeof StarlightFrontmatterSchema>;
+
+/** Plain object, union, and intersection Zod types. */
+type BaseSchemaWithoutEffects =
+	| z.AnyZodObject
+	| z.ZodUnion<[BaseSchemaWithoutEffects, ...BaseSchemaWithoutEffects[]]>
+	| z.ZodDiscriminatedUnion<string, z.AnyZodObject[]>
+	| z.ZodIntersection<BaseSchemaWithoutEffects, BaseSchemaWithoutEffects>;
+/** Base subset of Zod types that we support passing to the `extend` option. */
+type BaseSchema = BaseSchemaWithoutEffects | z.ZodEffects<BaseSchemaWithoutEffects>;
+
+/** Type that extends Starlight’s default schema with an optional, user-defined schema. */
+type ExtendedSchema<T extends BaseSchema> = T extends BaseSchema
+	? z.ZodIntersection<DefaultSchema, T>
+	: DefaultSchema;
+
+interface DocsSchemaOpts<T extends BaseSchema> {
+	/**
+	 * Extend Starlight’s schema with additional fields.
+	 *
+	 * @example
+	 * // Extend the built-in schema with a Zod schema.
+	 * docsSchema({
+	 * 	extend: z.object({
+	 * 		// Add a new field to the schema.
+	 * 		category: z.enum(['tutorial', 'guide', 'reference']).optional(),
+	 * 	}),
+	 * })
+	 *
+	 * // Use the Astro image helper.
+	 * docsSchema({
+	 * 	extend: ({ image }) => {
+	 * 		return z.object({
+	 * 			cover: image(),
+	 * 		});
+	 * 	},
+	 * })
+	 */
+	extend?: T | ((context: SchemaContext) => T);
+}
+
+/** Content collection schema for Starlight’s `docs` collection. */
+export function docsSchema<T extends BaseSchema>({ extend }: DocsSchemaOpts<T> = {}) {
+	return (context: SchemaContext): ExtendedSchema<T> => {
+		const UserSchema = typeof extend === 'function' ? extend(context) : extend;
+
+		return (
+			UserSchema
+				? StarlightFrontmatterSchema(context).and(UserSchema)
+				: StarlightFrontmatterSchema(context)
+		) as ExtendedSchema<T>;
+	};
 }
diff --git a/packages/starlight/schemas/i18n.ts b/packages/starlight/schemas/i18n.ts
index 0cf6496c..5a02e105 100644
--- a/packages/starlight/schemas/i18n.ts
+++ b/packages/starlight/schemas/i18n.ts
@@ -1,7 +1,31 @@
 import { z } from 'astro/zod';
 
-export function i18nSchema() {
-	return starlightI18nSchema().merge(pagefindI18nSchema()).merge(expressiveCodeI18nSchema());
+interface i18nSchemaOpts<T extends z.AnyZodObject = z.ZodObject<{}>> {
+	/**
+	 * Extend Starlight’s i18n schema with additional fields.
+	 *
+	 * @example
+	 * // Add two optional fields to the default schema.
+	 * i18nSchema({
+	 * 	extend: z
+	 * 		.object({
+	 * 			'customUi.heading': z.string(),
+	 * 			'customUi.text': z.string(),
+	 * 		})
+	 * 		.partial(),
+	 * })
+	 */
+	extend?: T;
+}
+
+/** Content collection schema for Starlight’s optional `i18n` collection. */
+export function i18nSchema<T extends z.AnyZodObject = z.ZodObject<{}>>({
+	extend = z.object({}) as T,
+}: i18nSchemaOpts<T> = {}) {
+	return starlightI18nSchema()
+		.merge(pagefindI18nSchema())
+		.merge(expressiveCodeI18nSchema())
+		.merge(extend);
 }
 export type i18nSchemaOutput = z.output<ReturnType<typeof i18nSchema>>;