From 62b0d20b974dc932769221d210b751627fb4bbc6 Mon Sep 17 00:00:00 2001
From: Martin Trapp <94928215+martrapp@users.noreply.github.com>
Date: Wed, 28 Aug 2024 13:09:49 +0200
Subject: [PATCH] Extends virtual module astro:transitions/client to export
 swapFunctions (#11708)

* extend virtual module astro:transitions/client to exports swapFunctions

* use virtual module in e2e tests

* Update .changeset/new-monkeys-sit.md

* Update .changeset/new-monkeys-sit.md

Co-authored-by: Sarah Rainsberger <sarah@rainsberger.ca>

* Update new-monkeys-sit.md

* Update swap-functions.ts

restoreFocus() bindings are now returned by saveFocus() and do not make sense anymore as a member of the swapFunctions object

* take over suggestion

Co-authored-by: Sarah Rainsberger <sarah@rainsberger.ca>

* Update .changeset/new-monkeys-sit.md

Co-authored-by: Sarah Rainsberger <sarah@rainsberger.ca>

---------

Co-authored-by: Emanuele Stoppa <my.burning@gmail.com>
Co-authored-by: Sarah Rainsberger <sarah@rainsberger.ca>
---
 .changeset/new-monkeys-sit.md                 | 32 +++++++++++++++++++
 packages/astro/client.d.ts                    |  2 ++
 .../src/pages/keep-style-one.astro            | 12 +++----
 .../src/pages/keep-theme-one.astro            | 12 +++----
 .../src/pages/replace-main-one.astro          | 14 ++++----
 packages/astro/e2e/view-transitions.test.js   |  2 +-
 .../astro/src/transitions/swap-functions.ts   |  8 +++++
 .../transitions/vite-plugin-transitions.ts    |  1 +
 .../transitions-swap-functions.ts             |  1 +
 9 files changed, 64 insertions(+), 20 deletions(-)
 create mode 100644 .changeset/new-monkeys-sit.md
 create mode 100644 packages/astro/src/virtual-modules/transitions-swap-functions.ts

diff --git a/.changeset/new-monkeys-sit.md b/.changeset/new-monkeys-sit.md
new file mode 100644
index 000000000000..b147b593f4a1
--- /dev/null
+++ b/.changeset/new-monkeys-sit.md
@@ -0,0 +1,32 @@
+---
+'astro': minor
+---
+
+Adds a new object `swapFunctions` to expose the necessary utility functions on `astro:transitions/client` that allow you to build custom swap functions to be used with view transitions. 
+
+The example below uses these functions to replace Astro's built-in default `swap` function with one that only swaps the `<main>` part of the page:
+
+```astro
+<script>
+import { swapFunctions } from 'astro:transitions/client';
+
+document.addEventListener('astro:before-swap', (e) => { e.swap = () => swapMainOnly(e.newDocument) });
+
+function swapMainOnly(doc: Document) {
+  swapFunctions.deselectScripts(doc);
+  swapFunctions.swapRootAttributes(doc);
+  swapFunctions.swapHeadElements(doc);
+  const restoreFocusFunction = swapFunctions.saveFocus();
+  const newMain = doc.querySelector('main');
+  const oldMain = document.querySelector('main');
+  if (newMain && oldMain) {
+    swapFunctions.swapBodyElement(newMain, oldMain);
+  } else {
+    swapFunctions.swapBodyElement(doc.body, document.body);
+  }
+  restoreFocusFunction();
+};
+<script>
+```
+
+See the [view transitions guide](https://docs.astro.build/en/guides/view-transitions/#astrobefore-swap) for more information about hooking into the `astro:before-swap` lifecycle event and adding a custom swap implementation.
diff --git a/packages/astro/client.d.ts b/packages/astro/client.d.ts
index ed5c1b894d28..11bb35cd783a 100644
--- a/packages/astro/client.d.ts
+++ b/packages/astro/client.d.ts
@@ -150,6 +150,8 @@ declare module 'astro:transitions/client' {
 		import('./dist/virtual-modules/transitions-events.js').TransitionBeforeSwapEvent;
 	export const isTransitionBeforePreparationEvent: EventModule['isTransitionBeforePreparationEvent'];
 	export const isTransitionBeforeSwapEvent: EventModule['isTransitionBeforeSwapEvent'];
+	type TransitionSwapFunctionModule = typeof import('./dist/virtual-modules/transitions-swap-functions.js');
+	export const swapFunctions: TransitionSwapFunctionModule['swapFunctions']
 }
 
 declare module 'astro:prefetch' {
diff --git a/packages/astro/e2e/fixtures/view-transitions/src/pages/keep-style-one.astro b/packages/astro/e2e/fixtures/view-transitions/src/pages/keep-style-one.astro
index 413a404d76fe..f71898d8153f 100644
--- a/packages/astro/e2e/fixtures/view-transitions/src/pages/keep-style-one.astro
+++ b/packages/astro/e2e/fixtures/view-transitions/src/pages/keep-style-one.astro
@@ -7,22 +7,22 @@ import Layout from '../components/Layout.astro';
 </Layout>
 <script>
 
-	import { deselectScripts, saveFocus, swapBodyElement, swapHeadElements, swapRootAttributes } from '../../node_modules/astro/dist/transitions/swap-functions'
+	import { swapFunctions } from 'astro:transitions/client';
 
 	document.addEventListener('astro:before-swap', (e) => {
 		e.swap = () => keepStyle(e.newDocument)
 	});
 
 	function keepStyle(doc: Document) {
-		deselectScripts(doc);
-		swapRootAttributes(doc);
+		swapFunctions.deselectScripts(doc);
+		swapFunctions.swapRootAttributes(doc);
 		{
 			const dynamicStyle = document.head.querySelector('style:not(:empty)');
-			swapHeadElements(doc);
+			swapFunctions.swapHeadElements(doc);
 			dynamicStyle && document.head.insertAdjacentElement('afterbegin', dynamicStyle);
 		}
-		const restoreFocusFunction = saveFocus();
-		swapBodyElement(doc.body, document.body)
+		const restoreFocusFunction = swapFunctions.saveFocus();
+		swapFunctions.swapBodyElement(doc.body, document.body)
 		restoreFocusFunction();
 	}
 </script>
diff --git a/packages/astro/e2e/fixtures/view-transitions/src/pages/keep-theme-one.astro b/packages/astro/e2e/fixtures/view-transitions/src/pages/keep-theme-one.astro
index a4c942d587f2..18ae0221fc65 100644
--- a/packages/astro/e2e/fixtures/view-transitions/src/pages/keep-theme-one.astro
+++ b/packages/astro/e2e/fixtures/view-transitions/src/pages/keep-theme-one.astro
@@ -6,18 +6,18 @@ import Layout from '../components/Layout.astro';
 	<a id="click" href="/keep-two">go to next page</a>
 </Layout>
 <script>
-	import { deselectScripts, saveFocus, swapBodyElement, swapHeadElements, swapRootAttributes } from '../../node_modules/astro/dist/transitions/swap-functions'
+	import { swapFunctions } from 'astro:transitions/client';
 
 	function keepTheme(doc:Document) {
-		deselectScripts(doc);
+		swapFunctions.deselectScripts(doc);
 		{
 			const theme = document.documentElement.getAttribute('data-theme')!;
-			swapRootAttributes(doc);
+			swapFunctions.swapRootAttributes(doc);
 			document.documentElement.setAttribute('data-theme', theme);
 		}
-		swapHeadElements(doc);
-		const restoreFocusFunction = saveFocus();
-		swapBodyElement(doc.body, document.body)
+		swapFunctions.swapHeadElements(doc);
+		const restoreFocusFunction = swapFunctions.saveFocus();
+		swapFunctions.swapBodyElement(doc.body, document.body)
 		restoreFocusFunction();
 	}
 
diff --git a/packages/astro/e2e/fixtures/view-transitions/src/pages/replace-main-one.astro b/packages/astro/e2e/fixtures/view-transitions/src/pages/replace-main-one.astro
index 1709661e8aca..c43ce78a9629 100644
--- a/packages/astro/e2e/fixtures/view-transitions/src/pages/replace-main-one.astro
+++ b/packages/astro/e2e/fixtures/view-transitions/src/pages/replace-main-one.astro
@@ -9,24 +9,24 @@ import Layout from '../components/Layout.astro';
 	<a id="click" href="/keep-two">go to next page</a>
 </Layout>
 <script>
-	import { deselectScripts, saveFocus, swapBodyElement, swapHeadElements, swapRootAttributes } from "../../node_modules/astro/dist/transitions/swap-functions"
+	import { swapFunctions } from 'astro:transitions/client';
 
 	document.addEventListener('astro:before-swap', (e) => {
 		e.swap = () => replaceMain(e.newDocument)
 	});
 
 	function replaceMain(doc:Document){
-		deselectScripts(doc);
-		swapRootAttributes(doc);
-		swapHeadElements(doc);
-		const restoreFocusFunction = saveFocus();
+		swapFunctions.deselectScripts(doc);
+		swapFunctions.swapRootAttributes(doc);
+		swapFunctions.swapHeadElements(doc);
+		const restoreFocusFunction = swapFunctions.saveFocus();
 		{
 			const newMain = doc.body.querySelector('main section');
 				const oldMain = document.body.querySelector('main section');
 				if (newMain && oldMain) {
-					swapBodyElement(newMain, oldMain);
+					swapFunctions.swapBodyElement(newMain, oldMain);
 				} else {
-					swapBodyElement(doc.body, document.body);
+					swapFunctions.swapBodyElement(doc.body, document.body);
 				}
 		}
 		restoreFocusFunction();
diff --git a/packages/astro/e2e/view-transitions.test.js b/packages/astro/e2e/view-transitions.test.js
index 822a57d89447..181994cfa930 100644
--- a/packages/astro/e2e/view-transitions.test.js
+++ b/packages/astro/e2e/view-transitions.test.js
@@ -1447,7 +1447,7 @@ test.describe('View Transitions', () => {
 		await page.click('#click');
 		await expect(page.locator('#name'), 'should have content').toHaveText('Keep 2');
 
-		const styleElement = await page.$('head > style');
+		const styleElement = await page.$('head > style:nth-child(1)');
 		const styleContent = await page.evaluate((style) => style.innerHTML, styleElement);
 		expect(styleContent).toBe('body { background-color: purple; }');
 	});
diff --git a/packages/astro/src/transitions/swap-functions.ts b/packages/astro/src/transitions/swap-functions.ts
index 4c8db82eef90..e2d8557f5d50 100644
--- a/packages/astro/src/transitions/swap-functions.ts
+++ b/packages/astro/src/transitions/swap-functions.ts
@@ -133,6 +133,14 @@ const shouldCopyProps = (el: HTMLElement): boolean => {
 	return persistProps == null || persistProps === 'false';
 };
 
+export const swapFunctions = {
+	deselectScripts,
+	swapRootAttributes,
+	swapHeadElements,
+	swapBodyElement,
+	saveFocus,
+};
+
 export const swap = (doc: Document) => {
 	deselectScripts(doc);
 	swapRootAttributes(doc);
diff --git a/packages/astro/src/transitions/vite-plugin-transitions.ts b/packages/astro/src/transitions/vite-plugin-transitions.ts
index d88c96f8919f..547ef3f38af0 100644
--- a/packages/astro/src/transitions/vite-plugin-transitions.ts
+++ b/packages/astro/src/transitions/vite-plugin-transitions.ts
@@ -42,6 +42,7 @@ export default function astroTransitions({ settings }: { settings: AstroSettings
 					TRANSITION_BEFORE_SWAP, isTransitionBeforeSwapEvent, TransitionBeforeSwapEvent,
 					TRANSITION_AFTER_SWAP, TRANSITION_PAGE_LOAD
 				} from "astro/virtual-modules/transitions-events.js";
+				export { swapFunctions } from "astro/virtual-modules/transitions-swap-functions.js";
 			`;
 			}
 		},
diff --git a/packages/astro/src/virtual-modules/transitions-swap-functions.ts b/packages/astro/src/virtual-modules/transitions-swap-functions.ts
new file mode 100644
index 000000000000..5947533e3d02
--- /dev/null
+++ b/packages/astro/src/virtual-modules/transitions-swap-functions.ts
@@ -0,0 +1 @@
+export * from '../transitions/swap-functions.js';