feat: lib support for streams and fixed headers

arthurfiorette · arthurfiorette · commit cdb5619c9c8f · 2023-09-29T18:02:30.000-03:00
diff --git a/src/handler.ts b/src/handler.ts
@@ -0,0 +1,68 @@
+import { Readable } from 'stream'
+import { HtmlOptions } from './options'
+import { isHtml, isTagHtml } from './utils'
+
+export function handleHtml(
+	value: string | Readable | Promise<string | Readable>,
+	options: HtmlOptions,
+	hasContentType: boolean
+): Promise<Response | string> | Response | string {
+	// Only use promises if value is a promise itself
+	if (value instanceof Promise) {
+		return value.then((v) => handleHtml(v, options, hasContentType))
+	}
+
+	// Simple string use cases
+	if (typeof value === 'string') {
+		if (
+			options.autoDoctype &&
+			isHtml(value) &&
+			// Avoids double adding !doctype or adding to non root html tags.
+			isTagHtml(value)
+		) {
+			value = '<!doctype html>' + value
+		}
+
+		return new Response(
+			value,
+			hasContentType
+				? undefined
+				: { headers: { 'content-type': options.contentType! } }
+		)
+	}
+
+	// Stream use cases
+	let stream = Readable.toWeb(value)
+
+	// We can convert to a readable stream with StreamTransform
+	if (options.autoDoctype) {
+		let first = true
+
+		stream = stream.pipeThrough(
+			new TransformStream({
+				transform(chunk, controller) {
+					let str = chunk!.toString()
+
+					if (
+						first &&
+						isTagHtml(str) &&
+						// Avoids double adding !doctype or adding to non root html tags.
+						isTagHtml(str)
+					) {
+						first = false
+						str = '<!doctype html>' + str
+					}
+
+					controller.enqueue(str)
+				}
+			})
+		)
+	}
+
+	return new Response(
+		stream,
+		hasContentType
+			? undefined
+			: { headers: { 'content-type': options.contentType! } }
+	)
+}
diff --git a/src/html.ts b/src/html.ts
@@ -1,6 +1,9 @@
 import { Elysia } from 'elysia'
-import { isHtml, isTagHtml } from './utils'
+import { Readable } from 'stream'
+import { handleHtml } from './handler'
 import { HtmlOptions } from './options'
+import { isHtml } from './utils'
+import { renderToStream } from '@kitajs/html/suspense'
 
 export function html(options: HtmlOptions = {}) {
 	// Defaults
@@ -9,48 +12,44 @@ export function html(options: HtmlOptions = {}) {
 	options.isHtml ??= isHtml
 	options.autoDoctype ??= true
 
-	let instance = new Elysia({ name: '@elysiajs/html' }).derive(({ set }) => ({
-		html(value: string) {
-			if (
-				options.autoDoctype &&
-				isHtml(value) &&
-				// Avoids double adding !doctype or adding to non root html tags.
-				isTagHtml(value)
-			) {
-				value = '<!doctype html>' + value
-			}
+	let instance = new Elysia({ name: '@elysiajs/html' }).derive(({ set }) => {
+		return {
+			html<A = void>(
+				value:
+					| Readable
+					| JSX.Element
+					| ((this: void, rid: number, ...args: A[]) => JSX.Element),
+				...args: A[]
+			): Promise<Response | string> | Response | string {
+				if (typeof value === 'function') {
+					value = renderToStream((rid) =>
+						(value as Function)(rid, ...args)
+					)
+				}
 
-			return new Response(value, {
-				...set,
-				// @ts-expect-error
-				headers: { ...set.headers, 'content-type': options.contentType! },
-			})
+				return handleHtml(value, options, 'content-type' in set.headers)
+			}
 		}
-	}))
+	})
 
 	if (options.autoDetect) {
-		instance = instance.onAfterHandle(
-			// onAfterHandle should be present on a lot of stack traces, so we should not
-			// use anonymous functions here.
-			function htmlHandle({ set }, response) {
-				if (!isHtml(response)) {
-					return response
-				}
-
-				// Full means that we should only try to convert raw string responses
-				if (options.autoDoctype === 'full' && isTagHtml(response)) {
-					response = '<!doctype html>' + response
-				}
-
-				set.headers['content-type'] = options.contentType!
-
-				return new Response(
-					// @ts-expect-error - We know this is a string.
-					response,
-					set
-				)
+		// handlerPossibleHtml should be present on a lot of stack traces, so we should not
+		// use anonymous functions here.
+		instance = instance.onAfterHandle(function handlerPossibleHtml({
+			response: value,
+			set
+		}) {
+			if (
+				// Simple html string
+				isHtml(value) ||
+				// @kitajs/html stream
+				(value instanceof Readable && 'rid' in value)
+			) {
+				return handleHtml(value, options, 'content-type' in set.headers)
 			}
-		)
+
+			return value
+		})
 	}
 
 	return instance
diff --git a/src/options.ts b/src/options.ts
@@ -21,16 +21,16 @@ export interface HtmlOptions {
 	/**
 	 * Whether to automatically add `<!doctype html>` to a response starting with <html>, if not found.
 	 *
-   * Use `full` to also automatically add doctypes on responses returned without this plugin
-   * 
-   * ```ts
-   * // without the plugin
-   * app.get('/', () => '<html></html>')
-   * 
-   * // With the plugin
-   * app.get('/', ({ html }) => html('<html></html>')	
-   * ```
-   * 
+	 * Use `full` to also automatically add doctypes on responses returned without this plugin
+	 *
+	 * ```ts
+	 * // without the plugin
+	 * app.get('/', () => '<html></html>')
+	 *
+	 * // With the plugin
+	 * app.get('/', ({ html }) => html('<html></html>')
+	 * ```
+	 *
 	 * @default true
 	 */
 	autoDoctype?: boolean | 'full'
diff --git a/src/utils.ts b/src/utils.ts
@@ -27,5 +27,5 @@ export function isHtml(this: void, value?: any): value is string {
  * casing**.
  */
 export function isTagHtml(this: void, value: string) {
-  return value.trimStart().slice(0, 5).startsWith('<html')
+	return value.trimStart().slice(0, 5).startsWith('<html')
 }

Original file line number	Diff line number	Diff line change
`@@ -27,5 +27,5 @@ export function isHtml(this: void, value?: any): value is string {`
`27`	`27`	`* casing**.`
`28`	`28`	`*/`
`29`	`29`	`export function isTagHtml(this: void, value: string) {`
`30`		`- return value.trimStart().slice(0, 5).startsWith('<html')`
	`30`	`+ return value.trimStart().slice(0, 5).startsWith('<html')`
`31`	`31`	`}`