Fix for augmentations

Author: AlemTuzlak

src/vite/plugin.tsx

@@ -115,31 +115,32 @@ export const reactRouterDevTools: (args?: ReactRouterViteConfig) => Plugin[] = (
 				return injectRdtClient(code, config, pluginImports, id)
 			},
 		},
+
 		{
-			name: "react-router-devtools-inject-context",
+			name: "react-router-devtools-data-function-augment",
 			apply(config) {
-				return shouldInject(config.mode, includeDevtools)
+				return shouldInject(config.mode, includeServer)
 			},
 			transform(code, id) {
 				const routeId = isTransformable(id)
 				if (!routeId) {
 					return
 				}
-				const finalCode = injectContext(code, routeId, id)
+				const finalCode = augmentDataFetchingFunctions(code, routeId, id)
 				return finalCode
 			},
 		},
 		{
-			name: "react-router-devtools-data-function-augment",
+			name: "react-router-devtools-inject-context",
 			apply(config) {
-				return shouldInject(config.mode, includeServer)
+				return shouldInject(config.mode, includeDevtools)
 			},
 			transform(code, id) {
 				const routeId = isTransformable(id)
 				if (!routeId) {
 					return
 				}
-				const finalCode = augmentDataFetchingFunctions(code, routeId, id)
+				const finalCode = injectContext(code, routeId, id)
 				return finalCode
 			},
 		},

src/vite/utils/data-functions-augment.test.ts

@@ -74,8 +74,9 @@ describe("transform", () => {
 			)
 			const expected = removeWhitespace(`
 			import { withLoaderWrapper as _withLoaderWrapper   } from "react-router-devtools/server";
-      export { loader as _loader } from "./loader.js";
+      import { loader as _loader } from "./loader.js";
 			export const loader = _withLoaderWrapper(_loader, "test");
+			export {} from "./loader.js";
 		`)
 			expect(removeWhitespace(result.code)).toStrictEqual(expected)
 		})
@@ -240,6 +241,24 @@ describe("transform", () => {
 			`)
 			expect(removeWhitespace(result.code)).toStrictEqual(expected)
 		})
+
+		it("should transform the client action export when it's re-exported from another file", () => {
+			const result = augmentDataFetchingFunctions(
+				`
+			export { clientLoader } from "./clientLoader.js";
+			`,
+				"test",
+				"/file/path"
+			)
+			const expected = removeWhitespace(`
+			import { withClientLoaderWrapper as _withClientLoaderWrapper   } from "react-router-devtools/client";
+      import { clientLoader as _clientLoader } from "./clientLoader.js";
+			export const clientLoader = _withClientLoaderWrapper(_clientLoader, "test");
+			export {} from "./clientLoader.js";
+		`)
+			expect(removeWhitespace(result.code)).toStrictEqual(expected)
+		})
+
 		it("should wrap the clientLoader export when it's exported via export { clientLoader } and declared within the file", () => {
 			const result = augmentDataFetchingFunctions(
 				`
@@ -338,8 +357,9 @@ describe("transform", () => {
 			)
 			const expected = removeWhitespace(`
 			import { withActionWrapper as _withActionWrapper   } from "react-router-devtools/server";
-      export { action as _action } from "./action.js";
+      import { action as _action } from "./action.js";
 			export const action = _withActionWrapper(_action, "test");
+			export {} from "./action.js";
 		`)
 			expect(removeWhitespace(result.code)).toStrictEqual(expected)
 		})
@@ -467,6 +487,23 @@ describe("transform", () => {
 			expect(removeWhitespace(result.code)).toStrictEqual(expected)
 		})
 
+		it("should transform the client action export when it's re-exported from another file", () => {
+			const result = augmentDataFetchingFunctions(
+				`
+			export { clientAction } from "./clientAction.js";
+			`,
+				"test",
+				"/file/path"
+			)
+			const expected = removeWhitespace(`
+			import { withClientActionWrapper as _withClientActionWrapper   } from "react-router-devtools/client";
+      import { clientAction as _clientAction } from "./clientAction.js";
+			export const clientAction = _withClientActionWrapper(_clientAction, "test");
+			export {} from "./clientAction.js";
+		`)
+			expect(removeWhitespace(result.code)).toStrictEqual(expected)
+		})
+
 		it("should transform the client action export when it's imported from another file and exported", () => {
 			const result = augmentDataFetchingFunctions(
 				`
@@ -512,3 +549,41 @@ describe("transform", () => {
 		})
 	})
 })
+
+it("should transform the re-exports when it's re-exported from another file with multiple re-exports", () => {
+	const result = augmentDataFetchingFunctions(
+		`
+	export { action, loader, default } from "./action.js";
+	`,
+		"test",
+		"/file/path"
+	)
+	const expected = removeWhitespace(`
+	import { withActionWrapper as _withActionWrapper, withLoaderWrapper as _withLoaderWrapper   } from "react-router-devtools/server";
+	import { action as _action } from "./action.js";
+	import { loader as _loader } from "./action.js";
+	export const action = _withActionWrapper(_action, "test");
+	export const loader = _withLoaderWrapper(_loader, "test");
+	export { default } from "./action.js";
+`)
+	expect(removeWhitespace(result.code)).toStrictEqual(expected)
+})
+
+it("should transform the re-exports when it's re-exported from another file with multiple re-exports", () => {
+	const result = augmentDataFetchingFunctions(
+		`
+	export { action, loader, default, blah } from "./action.js";
+	`,
+		"test",
+		"/file/path"
+	)
+	const expected = removeWhitespace(`
+	import { withActionWrapper as _withActionWrapper, withLoaderWrapper as _withLoaderWrapper   } from "react-router-devtools/server";
+	import { action as _action } from "./action.js";
+	import { loader as _loader } from "./action.js";
+	export const action = _withActionWrapper(_action, "test");
+	export const loader = _withLoaderWrapper(_loader, "test");
+	export { default, blah } from "./action.js";
+`)
+	expect(removeWhitespace(result.code)).toStrictEqual(expected)
+})

src/vite/utils/data-functions-augment.ts

@@ -32,6 +32,8 @@ const transform = (ast: ParseResult<Babel.File>, routeId: string) => {
 		return str.charAt(0).toUpperCase() + str.slice(1)
 	}
 	const transformations: Array<() => void> = []
+
+	const importDeclarations: Babel.ImportDeclaration[] = []
 	trav(ast, {
 		ExportDeclaration(path) {
 			if (path.isExportNamedDeclaration()) {
@@ -91,7 +93,39 @@ const transform = (ast: ParseResult<Babel.File>, routeId: string) => {
 					? getClientHocId(path, `with${uppercaseFirstLetter(name)}Wrapper`)
 					: getServerHocId(path, `with${uppercaseFirstLetter(name)}Wrapper`)
 				const binding = path.scope.getBinding(name)
-				if (binding?.path.isFunctionDeclaration()) {
+
+				if (path.node.source) {
+					// Special condition: export { loader, action } from "./path"
+					const source = path.node.source.value
+
+					importDeclarations.push(
+						t.importDeclaration(
+							[t.importSpecifier(t.identifier(`_${name}`), t.identifier(name))],
+							t.stringLiteral(source)
+						)
+					)
+					transformations.push(() => {
+						path.insertBefore(
+							t.exportNamedDeclaration(
+								t.variableDeclaration("const", [
+									t.variableDeclarator(
+										t.identifier(name),
+										t.callExpression(uid, [t.identifier(`_${name}`), t.stringLiteral(routeId)])
+									),
+								])
+							)
+						)
+					})
+
+					// Remove the specifier from the exports and add a manual export
+					transformations.push(() => {
+						const remainingSpecifiers = path.node.specifiers.filter(
+							(exportSpecifier) => !(t.isIdentifier(exportSpecifier.exported) && exportSpecifier.exported.name === name)
+						)
+
+						path.replaceWith(t.exportNamedDeclaration(null, remainingSpecifiers, path.node.source))
+					})
+				} else if (binding?.path.isFunctionDeclaration()) {
 					// Replace the function declaration with a wrapped version
 					binding.path.replaceWith(
 						t.variableDeclaration("const", [
@@ -138,6 +172,9 @@ const transform = (ast: ParseResult<Babel.File>, routeId: string) => {
 	for (const transformation of transformations) {
 		transformation()
 	}
+	if (importDeclarations.length > 0) {
+		ast.program.body.unshift(...importDeclarations)
+	}
 	if (serverHocs.length > 0) {
 		ast.program.body.unshift(
 			t.importDeclaration(

test-apps/react-router-vite/app/routes/exports.tsx

@@ -0,0 +1 @@
+export {  loader,  default  }from "./_index";