CHGenerator/Sources/Generator/Files/ImageGenerator.swift
2022-09-29 16:24:04 +02:00

288 lines
9.2 KiB
Swift

import Foundation
import AppKit
import CryptoKit
private struct ImageJob {
let destination: String
let width: Int
let path: String
}
final class ImageGenerator {
/**
The path to the input folder.
*/
private let input: URL
/**
The path to the output folder
*/
private let output: URL
/**
The images to generate.
The key is the image source path relative to the input folder, and the values are the destination path (relative to the output folder) and the required image width.
*/
private var imageJobs: [String : [ImageJob]] = [:]
/**
The images which could not be found, but are required for the site.
The key is the image path, and the value is the page that requires it.
*/
private var missingImages: [String : String] = [:]
/**
All warnings produced for images during generation
*/
private var imageWarnings: Set<String> = []
/**
All images required by the site.
The values are the destination paths of the images, relative to the output folder
*/
private var requiredImages: Set<String> = []
/**
All images modified or created during this generator run.
*/
private var generatedImages: Set<String> = []
/**
A cache to get the size of source images, so that files don't have to be loaded multiple times.
The key is the absolute source path, and the value is the image size
*/
private var imageSizeCache: [String : NSSize] = [:]
private var fileUpdates: FileUpdateChecker
init(input: URL, output: URL) {
self.fileUpdates = FileUpdateChecker(input: input)
self.input = input
self.output = output
}
func writeDetectedFileChangesToDisk() {
fileUpdates.writeDetectedFileChangesToDisk()
}
private func getImageSize(atPath path: String) -> NSSize? {
if let size = imageSizeCache[path] {
return size
}
guard let image = getImage(atPath: path) else {
return nil
}
let size = image.size
imageSizeCache[path] = size
return size
}
private func getImage(atPath path: String) -> NSImage? {
guard let data = getData(atPath: path) else {
log.add(error: "Failed to load file", source: path)
return nil
}
guard let image = NSImage(data: data) else {
log.add(error: "Failed to read image", source: path)
return nil
}
return image
}
private func getData(atPath path: String) -> Data? {
let url = input.appendingPathComponent(path)
guard url.exists else {
return nil
}
do {
let data = try Data(contentsOf: url)
fileUpdates.didLoad(data, at: path)
return data
} catch {
log.add(error: "Failed to read data", source: path, error: error)
return nil
}
}
func requireImage(at destination: String, generatedFrom source: String, requiredBy path: String, width: Int, height: Int?) -> NSSize {
requiredImages.insert(destination)
let height = height.unwrapped(CGFloat.init)
let sourceUrl = input.appendingPathComponent(source)
guard sourceUrl.exists else {
missingImages[source] = path
return .zero
}
guard let imageSize = getImageSize(atPath: source) else {
missingImages[source] = path
return .zero
}
let scaledSize = imageSize.scaledDown(to: CGFloat(width))
// Check desired height, then we can forget about it
if let height = height {
let expectedHeight = scaledSize.width / CGFloat(width) * height
if abs(expectedHeight - scaledSize.height) > 2 {
addWarning("Invalid height (\(scaledSize.height) instead of \(expectedHeight))", destination: destination, path: path)
}
}
let job = ImageJob(destination: destination, width: width, path: path)
guard let existingSource = imageJobs[source] else {
imageJobs[source] = [job]
return scaledSize
}
guard let existingJob = existingSource.first(where: { $0.destination == destination}) else {
imageJobs[source] = existingSource + [job]
return scaledSize
}
if existingJob.width != width {
addWarning("Multiple image widths (\(existingJob.width) and \(width))", destination: destination, path: "\(existingJob.path) and \(path)")
}
return scaledSize
}
func createImages() {
for (source, jobs) in imageJobs.sorted(by: { $0.key < $1.key }) {
create(images: jobs, from: source)
}
printMissingImages()
printImageWarnings()
printGeneratedImages()
printTotalImageCount()
}
private func printMissingImages() {
guard !missingImages.isEmpty else {
return
}
print("\(missingImages.count) missing images:")
for (source, path) in missingImages {
print(" \(source) (required by \(path))")
}
}
private func printImageWarnings() {
guard !imageWarnings.isEmpty else {
return
}
print("\(imageWarnings.count) image warnings:")
for imageWarning in imageWarnings {
print(imageWarning)
}
}
private func printGeneratedImages() {
guard !generatedImages.isEmpty else {
return
}
print("\(generatedImages.count) images generated:")
for image in generatedImages {
print(" " + image)
}
}
private func printTotalImageCount() {
print("\(requiredImages.count) images")
}
private func addWarning(_ message: String, destination: String, path: String) {
let warning = " \(destination): \(message) required by \(path)"
imageWarnings.insert(warning)
}
private func addWarning(_ message: String, job: ImageJob) {
addWarning(message, destination: job.destination, path: job.path)
}
private func isMissing(_ job: ImageJob) -> Bool {
!output.appendingPathComponent(job.destination).exists
}
private func create(images: [ImageJob], from source: String) {
// Only load image if required
let imageHasChanged = fileUpdates.fileHasChanged(at: source)
guard imageHasChanged || images.contains(where: isMissing) else {
return
}
guard let image = getImage(atPath: source) else {
missingImages[source] = images.first?.path
return
}
let jobs = imageHasChanged ? 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)
}
}
}
private func create(job: ImageJob, from image: NSImage, source: String) {
let destinationUrl = output.appendingPathComponent(job.destination)
// 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 {
addWarning("Invalid image extension \(destinationExtension)", job: job)
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(0.7)]) else {
addWarning("Failed to get data", job: job)
return
}
do {
try data.createFolderAndWrite(to: destinationUrl)
} catch {
addWarning("Failed to write image (\(error))", job: job)
return
}
generatedImages.insert(job.destination)
}
}