feat: add helper methods

Author: ascorbic

demos/middleware/middleware.ts

@@ -16,23 +16,10 @@ export async function middleware(req: NextRequest) {
     const res = await request.next()
     const message = `This was static but has been transformed in ${req.geo.city}`
 
-    // Transform the response page data
-    res.transformData((data) => {
-      data.pageProps.message = message
-      data.pageProps.showAd = true
-      return data
-    })
-
-    // Transform the response HTML
-    res.rewriteHTML('p[id=message]', {
-      text(textChunk) {
-        if (textChunk.lastInTextNode) {
-          textChunk.replace(message)
-        } else {
-          textChunk.remove()
-        }
-      },
-    })
+    // Transform the response HTML and props
+    res.replaceText('p[id=message]', message)
+    res.setPageProp('message', message)
+    res.setPageProp('showAd', true)
 
     return res
   }

demos/middleware/pages/static.js

@@ -1,18 +1,31 @@
-const Page = ({ message, showAd }) => (
-  <div>
-    <p id="message">{message}</p>
-    {showAd ? (
-      <div>
-        <p>This is an ad that isn't shown by default</p>
-        <img src="http://placekitten.com/400/300" />
-      </div>
-    ) : (
-      <p>No ads for me</p>
-    )}
-  </div>
-)
+import * as React from 'react'
 
-export async function getStaticProps(context) {
+const useHydrated = () => {
+  const [hydrated, setHydrated] = React.useState(false)
+  React.useEffect(() => {
+    setHydrated(true)
+  }, [])
+  return hydrated
+}
+
+const Page = ({ message, showAd }) => {
+  const hydrated = useHydrated()
+  return (
+    <div>
+      <p id="message">{message}</p>
+      {hydrated && showAd ? (
+        <div>
+          <p>This is an ad that isn't shown by default</p>
+          <img src="http://placekitten.com/400/300" />
+        </div>
+      ) : (
+        <p>No ads for me</p>
+      )}
+    </div>
+  )
+}
+
+export async function getStaticProps() {
   return {
     props: {
       message: 'This is a static page',

plugin/src/middleware/response.ts

@@ -3,7 +3,7 @@ import { NextResponse } from 'next/server'
 import type { ElementHandlers } from './html-rewriter'
 
 // eslint-disable-next-line @typescript-eslint/no-explicit-any
-export type NextDataTransform = <T extends Record<string, any>>(props: T) => T
+export type NextDataTransform = <T extends { pageProps?: Record<string, any> }>(props: T) => T
 
 // A NextReponse that wraps the Netlify origin response
 // We can't pass it through directly, because Next disallows returning a response body
@@ -42,6 +42,52 @@ export class MiddlewareResponse extends NextResponse {
     this.elementHandlers.push([selector, handlers])
   }
 
+  /**
+   * Sets the value of a page prop.
+   * @see transformData if you need more control
+   */
+  setPageProp(key: string, value: unknown) {
+    this.transformData((props) => {
+      props.pageProps ||= {}
+      props.pageProps[key] = value
+      return props
+    })
+  }
+
+  /**
+   * Replace the text of the given element. Takes either a string or a function
+   * that is passed the original string and returns new new string.
+   * @see rewriteHTML for more control
+   */
+  replaceText(selector: string, valueOrReplacer: string | ((input: string) => string)): void {
+    // If it's a string then our job is simpler, because we don't need to collect the current text
+    if (typeof valueOrReplacer === 'string') {
+      this.rewriteHTML(selector, {
+        text(textChunk) {
+          if (textChunk.lastInTextNode) {
+            textChunk.replace(valueOrReplacer)
+          } else {
+            textChunk.remove()
+          }
+        },
+      })
+    } else {
+      let text = ''
+      this.rewriteHTML(selector, {
+        text(textChunk) {
+          text += textChunk.text
+          // We're finished, so we can replace the text
+          if (textChunk.lastInTextNode) {
+            textChunk.replace(valueOrReplacer(text))
+          } else {
+            // Remove the chunk, because we'll be adding it back later
+            textChunk.remove()
+          }
+        },
+      })
+    }
+  }
+
   get headers(): Headers {
     // If we have the origin response, we should use its headers
     return this.originResponse?.headers || super.headers

plugin/src/templates/edge/utils.ts

@@ -73,10 +73,10 @@ export const buildResponse = async ({
     result.response = await result.response.next()
   }
   if (isMiddlewareResponse(result.response)) {
+    const { response } = result
     if (request.method === 'HEAD' || request.method === 'OPTIONS') {
       return response.originResponse
     }
-    const { response } = result
     // If it's JSON we don't need to use the rewriter, we can just parse it
     if (response.originResponse.headers.get('content-type')?.includes('application/json')) {
       const props = await response.originResponse.json()