type-erased HTMLResponse + headers (#4)

Author: sliemeobn

Package.swift

@@ -23,7 +23,7 @@ let package = Package(
     ],
     dependencies: [
         .package(url: "https://github.com/vapor/vapor.git", from: "4.102.0"),
-        .package(url: "https://github.com/sliemeobn/elementary.git", .upToNextMajor(from: "0.1.2")),
+        .package(url: "https://github.com/sliemeobn/elementary.git", .upToNextMajor(from: "0.3.0")),
     ],
     targets: [
         .target(

Sources/VaporElementary/HTMLResponse.swift

@@ -15,43 +15,53 @@ import Vapor
 ///   }
 /// }
 /// ```
-public struct HTMLResponse<Content: HTML & Sendable>: Sendable {
+public struct HTMLResponse: Sendable {
     // NOTE: The Sendable requirement on Content can probably be removed in Swift 6 using a sending parameter, and some fancy ~Copyable @unchecked Sendable box type.
     // We only need to pass the HTML value to the response generator body closure
-    private let content: Content
+    private let content: any HTML & Sendable
 
     /// The number of bytes to write to the response body at a time.
     ///
     /// The default is 1024 bytes.
     public var chunkSize: Int
 
+    /// Response headers
+    ///
+    /// It can be used to add additional headers to a predefined set of fields.
+    ///
+    /// - Note: If a new set of headers is assigned, all predefined headers are removed.
+    ///
+    /// ```swift
+    /// var response = HTMLResponse { ... }
+    /// response.headers.add(name: "foo", value: "bar")
+    /// return response
+    /// ```
+    public var headers: HTTPHeaders = ["Content-Type": "text/html; charset=utf-8"]
+
     /// Creates a new HTMLResponse
     ///
     /// - Parameters:
     ///   - chunkSize: The number of bytes to write to the response body at a time.
+    ///   - additionalHeaders: Additional headers to be merged with predefined headers.
     ///   - content: The `HTML` content to render in the response.
-    public init(chunkSize: Int = 1024, @HTMLBuilder content: () -> Content) {
+    public init(chunkSize: Int = 1024, additionalHeaders: HTTPHeaders = [:], @HTMLBuilder content: () -> some HTML & Sendable) {
         self.chunkSize = chunkSize
+        if additionalHeaders.contains(name: .contentType) {
+            self.headers = additionalHeaders
+        } else {
+            self.headers.add(contentsOf: additionalHeaders)
+        }
         self.content = content()
     }
 }
 
 extension HTMLResponse: AsyncResponseEncodable {
-    struct StreamWriter: HTMLStreamWriter {
-        var writer: any AsyncBodyStreamWriter
-        var allocator: ByteBufferAllocator
-
-        func write(_ bytes: ArraySlice<UInt8>) async throws {
-            try await self.writer.writeBuffer(self.allocator.buffer(bytes: bytes))
-        }
-    }
-
     public func encodeResponse(for request: Request) async throws -> Response {
         Response(
             status: .ok,
-            headers: ["Content-Type": "text/html; charset=utf-8"],
-            body: .init(asyncStream: { [content] writer in
-                try await content.render(into: StreamWriter(writer: writer, allocator: request.byteBufferAllocator))
+            headers: self.headers,
+            body: .init(asyncStream: { [content, chunkSize] writer in
+                try await writer.writeHTML(content, chunkSize: chunkSize)
                 try await writer.write(.end)
             })
         )

Sources/VaporElementary/HTMLResponseBodyWriter.swift

@@ -0,0 +1,25 @@
+import Elementary
+import Vapor
+
+struct HTMLResponseBodyStreamWriter: HTMLStreamWriter {
+    let allocator: ByteBufferAllocator = .init()
+    var writer: any AsyncBodyStreamWriter
+
+    mutating func write(_ bytes: ArraySlice<UInt8>) async throws {
+        try await self.writer.writeBuffer(self.allocator.buffer(bytes: bytes))
+    }
+}
+
+public extension AsyncBodyStreamWriter {
+    /// Writes HTML by rendering chuncks of bytes to the response body
+    ///
+    /// - Parameters:
+    ///   - html: The HTML content to render in the response
+    ///   - chunkSize: The number of bytes to write to the response body at a time (default is 1024 bytes)
+    func writeHTML(_ html: consuming some HTML, chunkSize: Int = 1204) async throws {
+        try await html.render(
+            into: HTMLResponseBodyStreamWriter(writer: self),
+            chunkSize: chunkSize
+        )
+    }
+}

Tests/VaporElementaryTests/HTMLResponseTests.swift

@@ -49,6 +49,47 @@ final class HTMLResponseTests: XCTestCase {
         let response = try await app.sendRequest(.GET, "/")
         XCTAssertEqual(String(buffer: response.body), Array(repeating: "<p></p>", count: count).joined())
     }
+
+    func testRespondsWithCustomHeaders() async throws {
+        self.app.get { _ in
+            var response = HTMLResponse(additionalHeaders: ["foo": "bar"]) { EmptyHTML() }
+            response.headers.add(name: "hx-refresh", value: "true")
+            return response
+        }
+
+        let response = try await app.sendRequest(.GET, "/")
+
+        XCTAssertEqual(response.headers["foo"], ["bar"])
+        XCTAssertEqual(response.headers["hx-refresh"], ["true"])
+        XCTAssertEqual(response.headers.contentType?.description, "text/html; charset=utf-8")
+    }
+
+    func testRespondsWithOverwrittenContentType() async throws {
+        self.app.get { _ in
+            HTMLResponse(additionalHeaders: ["Content-Type": "some"]) { EmptyHTML() }
+        }
+
+        let response = try await app.sendRequest(.GET, "/")
+
+        XCTAssertEqual(response.headers["Content-Type"], ["some"])
+    }
+
+    func testRespondsByWritingToStream() async throws {
+        self.app.get { _ in
+            Response(
+                status: .ok,
+                headers: [:],
+                body: .init(asyncStream: { writer in
+                    try await writer.writeHTML(p { "Hello" })
+                    try await writer.write(.end)
+                })
+            )
+        }
+
+        let response = try await app.sendRequest(.GET, "/")
+
+        XCTAssertEqual(String(buffer: response.body), "<p>Hello</p>")
+    }
 }
 
 struct TestPage: HTMLDocument {