CHGenerator/Sources/Generator/Processing/ImageGenerator.swift
2023-01-08 21:13:40 +01:00

347 lines
13 KiB
Swift

import Foundation
import AppKit
import CryptoKit
import Darwin.C
final class ImageGenerator {
private let imageOptimSupportedFileExtensions: Set<String> = ["jpg", "png", "svg"]
private let imageOptimizationBatchSize = 50
/**
The path to the input folder.
*/
private let input: URL
/**
The path to the output folder
*/
private let output: URL
private let imageReader: ImageReader
/**
All warnings produced for images during generation
*/
private var imageWarnings: [String]
/// The images which could not be found, but are required for the site (`key`: image path, `value`: source element path)
private var missingImages: [String : String]
/// Images which could not be read (`key`: file path relative to content, `value`: source element path)
private var unreadableImages: [String : String]
private var unhandledImages: [String: String] = [:]
/// All images modified or created during this generator run.
private var generatedImages: Set<String> = []
/// The images optimized by ImageOptim
private var optimizedImages: Set<String> = []
private var failedImages: [(path: String, message: String)] = []
private var numberOfGeneratedImages = 0
private let numberOfTotalImages: Int
private lazy var numberOfImagesToCreate = jobs.reduce(0) { $0 + $1.images.count } + multiJobs.count * 2
private var numberOfImagesToOptimize = 0
private var numberOfOptimizedImages = 0
private let images: ImageData
private lazy var jobs: [(source: String, images: [ImageJob])] = images.jobs
.sorted { $0.key < $1.key }
.map { (source: $0.key, images: $0.value) }
.filter {
// Only load image if required
let imageHasChanged = imageReader.imageHasChanged(at: $0.source)
return imageHasChanged || $0.images.contains { job in
job.alwaysGenerate || !output.appendingPathComponent(job.destination).exists
}
}
private lazy var multiJobs: [String : String] = {
let imagesToGenerate: Set<String> = jobs.reduce([]) { $0.union($1.images.map { $0.destination }) }
return images.multiJobs.filter { imagesToGenerate.contains($0.key) }
}()
init(input: URL, output: URL, reader: ImageReader, images: ImageData) {
self.input = input
self.output = output
self.imageReader = reader
self.images = images
self.numberOfTotalImages = images.jobs.reduce(0) { $0 + $1.value.count }
+ images.multiJobs.count * 2
self.imageWarnings = images.warnings
self.missingImages = images.missing
self.unreadableImages = images.unreadable
}
func generateImages() {
var notes: [String] = []
func addIfNotZero(_ count: Int, _ name: String) {
guard count > 0 else {
return
}
notes.append("\(count) \(name)")
}
print(" Changed sources: \(jobs.count)/\(images.jobs.count)")
print(" Total images: \(numberOfTotalImages) (\(numberOfTotalImages - imageReader.numberOfFilesAccessed) versions)")
for (source, jobs) in jobs {
create(images: jobs, from: source)
}
for (baseImage, source) in multiJobs {
createMultiImages(from: baseImage, path: source)
}
print(" Generated images: \(numberOfGeneratedImages)/\(numberOfImagesToCreate)")
optimizeImages()
print(" Optimized images: \(numberOfOptimizedImages)/\(numberOfImagesToOptimize)")
addIfNotZero(missingImages.count, "missing images")
addIfNotZero(unreadableImages.count, "unreadable images")
print(" Warnings: \(imageWarnings.count)")
if !notes.isEmpty {
print(" Notes: " + notes.joined(separator: ", "))
}
}
private func create(images: [ImageJob], from source: String) {
guard let image = imageReader.getImage(atPath: source) else {
unreadableImages[source] = images.first!.destination
didGenerateImage(count: images.count)
return
}
let jobs = imageReader.imageHasChanged(at: source) ? images : images.filter(isMissing)
// Update all images
jobs.forEach { job in
// Prevent memory overflow due to repeated NSImage operations
autoreleasepool {
create(job: job, from: image, source: source)
didGenerateImage()
}
}
}
private func isMissing(_ job: ImageJob) -> Bool {
job.alwaysGenerate || !output.appendingPathComponent(job.destination).exists
}
private func create(job: ImageJob, from image: NSImage, source: String) {
let destinationUrl = output.appendingPathComponent(job.destination)
create(job: job, from: image, source: source, at: destinationUrl)
}
private func create(job: ImageJob, from image: NSImage, source: String, at destinationUrl: URL) {
// Ensure that image file is supported
let ext = destinationUrl.pathExtension.lowercased()
guard ImageType(fileExtension: ext) != nil else {
fatalError()
}
let destinationExtension = destinationUrl.pathExtension.lowercased()
guard let type = ImageType(fileExtension: destinationExtension)?.fileType else {
unhandledImages[source] = job.destination
return
}
let desiredWidth = CGFloat(job.width)
let sourceRep = image.representations[0]
let destinationSize = NSSize(width: sourceRep.pixelsWide, height: sourceRep.pixelsHigh)
.scaledDown(to: desiredWidth)
// create NSBitmapRep manually, if using cgImage, the resulting size is wrong
let rep = NSBitmapImageRep(bitmapDataPlanes: nil,
pixelsWide: Int(destinationSize.width),
pixelsHigh: Int(destinationSize.height),
bitsPerSample: 8,
samplesPerPixel: 4,
hasAlpha: true,
isPlanar: false,
colorSpaceName: NSColorSpaceName.deviceRGB,
bytesPerRow: Int(destinationSize.width) * 4,
bitsPerPixel: 32)!
let ctx = NSGraphicsContext(bitmapImageRep: rep)
NSGraphicsContext.saveGraphicsState()
NSGraphicsContext.current = ctx
image.draw(in: NSMakeRect(0, 0, destinationSize.width, destinationSize.height))
ctx?.flushGraphics()
NSGraphicsContext.restoreGraphicsState()
// Get NSData, and save it
guard let data = rep.representation(using: type, properties: [.compressionFactor: NSNumber(value: job.quality)]) else {
markImageAsFailed(source, error: "Failed to get data")
return
}
do {
try data.createFolderAndWrite(to: destinationUrl)
} catch {
markImageAsFailed(job.destination, error: "Failed to write image (\(error))")
return
}
generatedImages.insert(job.destination)
}
private func markImageAsFailed(_ source: String, error: String) {
failedImages.append((source, error))
}
private func createMultiImages(from source: String, path: String) {
guard generatedImages.contains(source) else {
didGenerateImage(count: 2)
return
}
let sourceUrl = output.appendingPathComponent(source)
let sourcePath = sourceUrl.path
guard sourceUrl.exists else {
missingImages[source] = path
didGenerateImage(count: 2)
return
}
let avifPath = source.dropAfterLast(".") + ".avif"
createAVIF(at: output.appendingPathComponent(avifPath).path, from: sourcePath)
generatedImages.insert(avifPath)
didGenerateImage()
let webpPath = source.dropAfterLast(".") + ".webp"
createWEBP(at: output.appendingPathComponent(webpPath).path, from: sourcePath)
generatedImages.insert(webpPath)
didGenerateImage()
compress(at: sourcePath)
}
private func createAVIF(at destination: String, from source: String, quality: Int = 55, effort: Int = 5) {
let folder = destination.dropAfterLast("/")
let command = "npx avif --input=\(source) --quality=\(quality) --effort=\(effort) --output=\(folder) --overwrite"
do {
let output = try safeShell(command)
if output == "" {
return
}
markImageAsFailed(destination, error: "Failed to create AVIF image: \(output)")
} catch {
markImageAsFailed(destination, error: "Failed to create AVIF image")
}
}
private func createWEBP(at destination: String, from source: String, quality: Int = 75) {
let command = "cwebp \(source) -q \(quality) -o \(destination)"
do {
let output = try safeShell(command)
if !output.contains("Error") {
return
}
markImageAsFailed(destination, error: "Failed to create WEBP image: \(output)")
} catch {
markImageAsFailed(destination, error: "Failed to create WEBP image: \(error)")
}
}
private func compress(at destination: String, quality: Int = 70) {
let command = "magick convert \(destination) -quality \(quality)% \(destination)"
do {
let output = try safeShell(command)
if output == "" {
return
}
markImageAsFailed(destination, error: "Failed to compress image: \(output)")
} catch {
markImageAsFailed(destination, error: "Failed to compress image: \(error)")
}
}
private func optimizeImages() {
let all = generatedImages
.filter { imageOptimSupportedFileExtensions.contains($0.lastComponentAfter(".")) }
.map { output.appendingPathComponent($0).path }
numberOfImagesToOptimize = all.count
for i in stride(from: 0, to: numberOfImagesToOptimize, by: imageOptimizationBatchSize) {
let endIndex = min(i+imageOptimizationBatchSize, numberOfImagesToOptimize)
let batch = all[i..<endIndex]
if optimizeImageBatch(batch) {
optimizedImages.formUnion(batch)
}
didOptimizeImage(count: batch.count)
}
}
private func optimizeImageBatch(_ batch: ArraySlice<String>) -> Bool {
let command = "imageoptim " + batch.joined(separator: " ")
do {
let output = try safeShell(command)
if output.contains("Finished") {
return true
}
for image in batch {
markImageAsFailed(image, error: "Failed to optimize image: \(output)")
}
return true
} catch {
for image in batch {
markImageAsFailed(image, error: "Failed to optimize image: \(error)")
}
return false
}
}
// MARK: Output
private func didGenerateImage(count: Int = 1) {
numberOfGeneratedImages += count
print(" Generated images: \(numberOfGeneratedImages)/\(numberOfImagesToCreate) \r", terminator: "")
fflush(stdout)
}
private func didOptimizeImage(count: Int) {
numberOfOptimizedImages += count
print(" Optimized images: \(numberOfOptimizedImages)/\(numberOfImagesToOptimize) \r", terminator: "")
fflush(stdout)
}
func writeResults(to file: URL) {
guard !missingImages.isEmpty || !unreadableImages.isEmpty || !failedImages.isEmpty || !unhandledImages.isEmpty || !imageWarnings.isEmpty || !generatedImages.isEmpty || !optimizedImages.isEmpty else {
do {
try FileManager.default.removeItem(at: file)
} catch {
print(" Failed to delete image log: \(error)")
}
return
}
var lines: [String] = []
func add<S>(_ name: String, _ property: S, convert: (S.Element) -> String) where S: Sequence {
let elements = property.map { " " + convert($0) }.sorted()
guard !elements.isEmpty else {
return
}
lines.append("\(name):")
lines.append(contentsOf: elements)
}
add("Missing images", missingImages) { "\($0.key) (required by \($0.value))" }
add("Unreadable images", unreadableImages) { "\($0.key) (required by \($0.value))" }
add("Failed images", failedImages) { "\($0.path): \($0.message)" }
add("Unhandled images", unhandledImages) { "\($0.value) (from \($0.key))" }
add("Warnings", imageWarnings) { $0 }
add("Generated images", generatedImages) { $0 }
add("Optimized images", optimizedImages) { $0 }
let data = lines.joined(separator: "\n").data(using: .utf8)!
do {
try data.createFolderAndWrite(to: file)
} catch {
print(" Failed to save image log: \(error)")
}
}
}