├── Tests └── ocritTests │ ├── Resources │ ├── test-en.png │ ├── test-pt.png │ ├── test-zh.png │ ├── ocrit-fixtures.sketch │ ├── test-en-multipage.pdf │ ├── test-multi-en-ko.png │ └── test-en-singlepage.pdf │ ├── StringHelpers.swift │ ├── ocritTests.swift │ └── TestHelpers.swift ├── Sources └── ocrit │ ├── Path+ArgumentParser.swift │ ├── Implementation │ ├── Base │ │ ├── OCROperation.swift │ │ ├── TranslationOperation.swift │ │ ├── VNRecognizeTextRequest+Validation.swift │ │ └── CGImageOCR.swift │ ├── CGPDFDocument+CGImage.swift │ ├── ImageOCROperation.swift │ ├── PDFOCROperation.swift │ └── AppleTranslateOperation.swift │ ├── Output.swift │ └── OCRIT.swift ├── .gitignore ├── Package.resolved ├── Package.swift ├── LICENSE └── README.md /Tests/ocritTests/Resources/test-en.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/insidegui/ocrit/HEAD/Tests/ocritTests/Resources/test-en.png -------------------------------------------------------------------------------- /Tests/ocritTests/Resources/test-pt.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/insidegui/ocrit/HEAD/Tests/ocritTests/Resources/test-pt.png -------------------------------------------------------------------------------- /Tests/ocritTests/Resources/test-zh.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/insidegui/ocrit/HEAD/Tests/ocritTests/Resources/test-zh.png -------------------------------------------------------------------------------- /Tests/ocritTests/Resources/ocrit-fixtures.sketch: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/insidegui/ocrit/HEAD/Tests/ocritTests/Resources/ocrit-fixtures.sketch -------------------------------------------------------------------------------- /Tests/ocritTests/Resources/test-en-multipage.pdf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/insidegui/ocrit/HEAD/Tests/ocritTests/Resources/test-en-multipage.pdf -------------------------------------------------------------------------------- /Tests/ocritTests/Resources/test-multi-en-ko.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/insidegui/ocrit/HEAD/Tests/ocritTests/Resources/test-multi-en-ko.png -------------------------------------------------------------------------------- /Tests/ocritTests/Resources/test-en-singlepage.pdf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/insidegui/ocrit/HEAD/Tests/ocritTests/Resources/test-en-singlepage.pdf -------------------------------------------------------------------------------- /Sources/ocrit/Path+ArgumentParser.swift: -------------------------------------------------------------------------------- 1 | import ArgumentParser 2 | import PathKit 3 | 4 | extension Path: @retroactive ExpressibleByArgument { 5 | public init?(argument: String) { 6 | self = Path(argument).absolute() 7 | } 8 | } 9 | -------------------------------------------------------------------------------- /Sources/ocrit/Implementation/Base/OCROperation.swift: -------------------------------------------------------------------------------- 1 | import Foundation 2 | 3 | struct OCRResult { 4 | var text: String 5 | var suggestedFilename: String 6 | } 7 | 8 | protocol OCROperation { 9 | init(fileURL: URL, customLanguages: [String]) 10 | func run(fast: Bool) throws -> AsyncThrowingStream 11 | } 12 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | .DS_Store 2 | __MACOSX 3 | *.pbxuser 4 | !default.pbxuser 5 | *.mode1v3 6 | !default.mode1v3 7 | *.mode2v3 8 | !default.mode2v3 9 | *.perspectivev3 10 | !default.perspectivev3 11 | *.xcworkspace 12 | !default.xcworkspace 13 | xcuserdata 14 | profile 15 | *.moved-aside 16 | DerivedData 17 | .idea/ 18 | Crashlytics.sh 19 | generatechangelog.sh 20 | Pods/ 21 | Carthage 22 | Provisioning 23 | Crashlytics.sh 24 | Sharing.h 25 | Tests/Private 26 | Design/Icon 27 | Build/ 28 | MockServer/ 29 | .swiftpm -------------------------------------------------------------------------------- /Sources/ocrit/Output.swift: -------------------------------------------------------------------------------- 1 | import ArgumentParser 2 | import Foundation 3 | import PathKit 4 | 5 | enum Output { 6 | case stdOutput 7 | case path(Path) 8 | } 9 | 10 | extension Output: ExpressibleByArgument { 11 | init?(argument: String) { 12 | if argument == "-" { 13 | self = .stdOutput 14 | } 15 | 16 | let path = Path(argument).absolute() 17 | self = .path(path) 18 | } 19 | } 20 | 21 | extension Output { 22 | var isStdOutput: Bool { 23 | switch self { 24 | case .stdOutput: true 25 | default: false 26 | } 27 | } 28 | 29 | var path: Path? { 30 | switch self { 31 | case let .path(path): path 32 | default: nil 33 | } 34 | } 35 | } 36 | -------------------------------------------------------------------------------- /Sources/ocrit/Implementation/Base/TranslationOperation.swift: -------------------------------------------------------------------------------- 1 | import Foundation 2 | 3 | struct TranslationResult { 4 | var sourceText: String 5 | var translatedText: String 6 | var inputLanguage: String 7 | var outputLanguage: String 8 | } 9 | 10 | enum TranslationAvailability { 11 | /// Translation is supported, but one or more languages are not installed. 12 | case supported 13 | /// Translation is supported and all required languages are installed. 14 | case installed 15 | /// Translation is not supported. 16 | case unsupported 17 | } 18 | 19 | protocol TranslationOperation { 20 | static func availability(from inputLanguage: String, to outputLanguage: String) async -> TranslationAvailability 21 | init(text: String, inputLanguage: String, outputLanguage: String) 22 | func run() async throws -> TranslationResult 23 | } 24 | -------------------------------------------------------------------------------- /Tests/ocritTests/StringHelpers.swift: -------------------------------------------------------------------------------- 1 | //===----------------------------------------------------------*- swift -*-===// 2 | // 3 | // This source file is part of the Swift Argument Parser open source project 4 | // 5 | // Copyright (c) 2020 Apple Inc. and the Swift project authors 6 | // Licensed under Apache License v2.0 with Runtime Library Exception 7 | // 8 | // See https://swift.org/LICENSE.txt for license information 9 | // 10 | //===----------------------------------------------------------------------===// 11 | 12 | extension Substring { 13 | func trimmed() -> Substring { 14 | guard let i = lastIndex(where: { $0 != " "}) else { 15 | return "" 16 | } 17 | return self[...i] 18 | } 19 | } 20 | 21 | extension String { 22 | public func trimmingLines() -> String { 23 | return self 24 | .split(separator: "\n", omittingEmptySubsequences: false) 25 | .map { $0.trimmed() } 26 | .joined(separator: "\n") 27 | } 28 | } 29 | -------------------------------------------------------------------------------- /Package.resolved: -------------------------------------------------------------------------------- 1 | { 2 | "pins" : [ 3 | { 4 | "identity" : "pathkit", 5 | "kind" : "remoteSourceControl", 6 | "location" : "https://github.com/kylef/PathKit", 7 | "state" : { 8 | "revision" : "3bfd2737b700b9a36565a8c94f4ad2b050a5e574", 9 | "version" : "1.0.1" 10 | } 11 | }, 12 | { 13 | "identity" : "spectre", 14 | "kind" : "remoteSourceControl", 15 | "location" : "https://github.com/kylef/Spectre.git", 16 | "state" : { 17 | "revision" : "26cc5e9ae0947092c7139ef7ba612e34646086c7", 18 | "version" : "0.10.1" 19 | } 20 | }, 21 | { 22 | "identity" : "swift-argument-parser", 23 | "kind" : "remoteSourceControl", 24 | "location" : "https://github.com/apple/swift-argument-parser.git", 25 | "state" : { 26 | "revision" : "f3c9084a71ef4376f2fabbdf1d3d90a49f1fabdb", 27 | "version" : "1.1.2" 28 | } 29 | } 30 | ], 31 | "version" : 2 32 | } 33 | -------------------------------------------------------------------------------- /Sources/ocrit/Implementation/Base/VNRecognizeTextRequest+Validation.swift: -------------------------------------------------------------------------------- 1 | import ArgumentParser 2 | import Vision 3 | 4 | extension VNRecognizeTextRequest { 5 | @discardableResult 6 | static func validateLanguages(with customLanguages: [String]) throws -> [String]? { 7 | let dummy = VNRecognizeTextRequest() 8 | return try dummy.validateLanguages(with: customLanguages) 9 | } 10 | 11 | func validateLanguages(with customLanguages: [String]) throws -> [String]? { 12 | guard !customLanguages.isEmpty else { return nil } 13 | 14 | let supportedLanguages = try supportedRecognitionLanguages() 15 | 16 | for customLanguage in customLanguages { 17 | guard supportedLanguages.contains(customLanguage) else { 18 | throw ValidationError("Unsupported language \"\(customLanguage)\". Supported languages are: \(supportedLanguages.joined(separator: ", "))") 19 | } 20 | } 21 | 22 | return customLanguages 23 | } 24 | } 25 | -------------------------------------------------------------------------------- /Sources/ocrit/Implementation/CGPDFDocument+CGImage.swift: -------------------------------------------------------------------------------- 1 | import Quartz 2 | 3 | extension CGPDFDocument { 4 | func cgImage(at pageNumber: Int) throws -> CGImage { 5 | guard let page = page(at: pageNumber) else { 6 | throw Failure("Page #\(pageNumber) not found.") 7 | } 8 | 9 | let pageRect = page.getBoxRect(.mediaBox) 10 | 11 | let img = NSImage(size: pageRect.size, flipped: true) { rect in 12 | guard let ctx = NSGraphicsContext.current?.cgContext else { return false } 13 | 14 | NSColor.white.setFill() 15 | rect.fill() 16 | 17 | ctx.translateBy(x: 0, y: pageRect.size.height) 18 | ctx.scaleBy(x: 1.0, y: -1.0) 19 | 20 | ctx.drawPDFPage(page) 21 | 22 | return true 23 | } 24 | 25 | guard let cgImage = img.cgImage(forProposedRect: nil, context: nil, hints: nil) else { 26 | throw Failure("Failed to create CGImage.") 27 | } 28 | 29 | return cgImage 30 | } 31 | } 32 | -------------------------------------------------------------------------------- /Package.swift: -------------------------------------------------------------------------------- 1 | // swift-tools-version: 5.6 2 | 3 | import PackageDescription 4 | 5 | let package = Package( 6 | name: "ocrit", 7 | platforms: [.macOS(.v12)], 8 | dependencies: [ 9 | .package(url: "https://github.com/apple/swift-argument-parser.git", from: "1.1.0"), 10 | .package(url: "https://github.com/kylef/PathKit", from: "1.0.1") 11 | ], 12 | targets: [ 13 | .executableTarget( 14 | name: "ocrit", 15 | dependencies: [ 16 | .product(name: "ArgumentParser", package: "swift-argument-parser"), 17 | .product(name: "PathKit", package: "PathKit") 18 | ]), 19 | .testTarget( 20 | name: "ocritTests", 21 | dependencies: ["ocrit"], 22 | exclude: ["Resources/ocrit-fixtures.sketch"], 23 | resources: [ 24 | .copy("Resources/test-en.png"), 25 | .copy("Resources/test-pt.png"), 26 | .copy("Resources/test-zh.png"), 27 | .copy("Resources/test-multi-en-ko.png"), 28 | .copy("Resources/test-en-singlepage.pdf"), 29 | .copy("Resources/test-en-multipage.pdf"), 30 | ] 31 | ), 32 | ] 33 | ) 34 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | Copyright (c) 2022 Guilherme Rambo 2 | 3 | Redistribution and use in source and binary forms, with or without 4 | modification, are permitted provided that the following conditions are met: 5 | 6 | - Redistributions of source code must retain the above copyright notice, this 7 | list of conditions and the following disclaimer. 8 | 9 | - Redistributions in binary form must reproduce the above copyright notice, 10 | this list of conditions and the following disclaimer in the documentation 11 | and/or other materials provided with the distribution. 12 | 13 | THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" 14 | AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 15 | IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE 16 | DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE 17 | FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL 18 | DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR 19 | SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER 20 | CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, 21 | OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 22 | OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. -------------------------------------------------------------------------------- /Sources/ocrit/Implementation/ImageOCROperation.swift: -------------------------------------------------------------------------------- 1 | import Vision 2 | import Cocoa 3 | 4 | final class ImageOCROperation: OCROperation { 5 | 6 | let imageURL: URL 7 | let customLanguages: [String] 8 | 9 | init(fileURL: URL, customLanguages: [String]) { 10 | self.imageURL = fileURL 11 | self.customLanguages = customLanguages 12 | } 13 | 14 | func run(fast: Bool) throws -> AsyncThrowingStream { 15 | guard let image = NSImage(contentsOf: imageURL) else { 16 | throw Failure("Couldn't read image at \(imageURL.path)") 17 | } 18 | 19 | guard let cgImage = image.cgImage(forProposedRect: nil, context: nil, hints: nil) else { 20 | throw Failure("Couldn't read CGImage fir \(imageURL.lastPathComponent)") 21 | } 22 | 23 | let filename = imageURL.deletingPathExtension().lastPathComponent 24 | 25 | let ocr = CGImageOCR(image: cgImage, customLanguages: customLanguages) 26 | 27 | return AsyncThrowingStream { continuation in 28 | Task { 29 | do { 30 | let text = try await ocr.run(fast: fast) 31 | 32 | let result = OCRResult(text: text, suggestedFilename: filename) 33 | 34 | continuation.yield(result) 35 | continuation.finish() 36 | } catch { 37 | continuation.finish(throwing: error) 38 | } 39 | } 40 | } 41 | } 42 | 43 | } 44 | -------------------------------------------------------------------------------- /Sources/ocrit/Implementation/PDFOCROperation.swift: -------------------------------------------------------------------------------- 1 | import Vision 2 | import Quartz 3 | 4 | final class PDFOCROperation: OCROperation { 5 | 6 | let documentURL: URL 7 | let customLanguages: [String] 8 | 9 | init(fileURL: URL, customLanguages: [String]) { 10 | self.documentURL = fileURL 11 | self.customLanguages = customLanguages 12 | } 13 | 14 | func run(fast: Bool) throws -> AsyncThrowingStream { 15 | let basename = documentURL.deletingPathExtension().lastPathComponent 16 | 17 | guard let document = CGPDFDocument(documentURL as CFURL) else { 18 | throw Failure("Failed to read PDF at \(documentURL.path)") 19 | } 20 | 21 | guard document.numberOfPages > 0 else { 22 | throw Failure("PDF has no pages at \(documentURL.path)") 23 | } 24 | 25 | return AsyncThrowingStream { continuation in 26 | Task { 27 | for page in (1...document.numberOfPages) { 28 | do { 29 | let cgImage = try document.cgImage(at: page) 30 | 31 | let ocr = CGImageOCR(image: cgImage, customLanguages: customLanguages) 32 | 33 | let text = try await ocr.run(fast: fast) 34 | 35 | let result = OCRResult(text: text, suggestedFilename: basename + "-\(page)") 36 | 37 | continuation.yield(result) 38 | } catch { 39 | /// Don't want to interrupt processing if a single page fails, so don't terminate the stream here. 40 | fputs("WARN: Error processing PDF page #\(page) at \(documentURL.path): \(error)\n", stderr) 41 | } 42 | } 43 | 44 | continuation.finish() 45 | } 46 | } 47 | } 48 | 49 | } 50 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # ocrit 2 | 3 | Runs Vision's OCR on input images or PDF files and outputs corresponding `txt` files for each image, or writes the recognized results to standard output. 4 | 5 | ``` 6 | USAGE: ocrit [ ...] [--output ] [--language ...] [--fast] 7 | 8 | ARGUMENTS: 9 | Path or list of paths for the images 10 | 11 | OPTIONS: 12 | -o, --output Path to a directory where the txt files will be written to, or - for standard output (default: -) 13 | -l, --language 14 | Language code to use for the recognition, can be repeated to select multiple languages 15 | -f, --fast Uses an OCR algorithm that prioritizes speed over accuracy 16 | -h, --help Show help information. 17 | ``` 18 | 19 | ## Language Selection 20 | 21 | The `--language` (or `-l`) option can be used to indicate which language or languages will be used for OCR. 22 | 23 | Multiple languages can be specified by repeating the option, example: 24 | 25 | ``` 26 | ocrit path/to/image.png -l ko-KR -l en-US 27 | ``` 28 | 29 | The order of the languages is important, as Vision's OCR engine will attempt to perform OCR using the languages in order. In my experience, if you have an image or document that contains a mix of English and some other language, it's best to specify `en-US` as the **last** language on the list. 30 | 31 | ### Supported Languages 32 | 33 | Language support varies with the version of macOS and whether or not the `--fast` flag is specified. 34 | 35 | This is the current list of supported languages as of macOS 14.4: 36 | 37 | ``` 38 | en-US, fr-FR, it-IT, de-DE, es-ES, pt-BR, zh-Hans, zh-Hant, yue-Hans, yue-Hant, ko-KR, ja-JP, ru-RU, uk-UA, th-TH, vi-VT 39 | ``` 40 | 41 | This is the current list of supported languages as of macOS 14.4, with the `--fast` flag enabled: 42 | 43 | ``` 44 | en-US, fr-FR, it-IT, de-DE, es-ES, pt-BR 45 | ``` -------------------------------------------------------------------------------- /Sources/ocrit/Implementation/Base/CGImageOCR.swift: -------------------------------------------------------------------------------- 1 | import Vision 2 | import Cocoa 3 | 4 | final class CGImageOCR { 5 | 6 | let image: CGImage 7 | let customLanguages: [String] 8 | 9 | init(image: CGImage, customLanguages: [String]) { 10 | self.image = image 11 | self.customLanguages = customLanguages 12 | } 13 | 14 | private var request: VNRecognizeTextRequest? 15 | private var handler: VNImageRequestHandler? 16 | 17 | func run(fast: Bool) async throws -> String { 18 | return try await withCheckedThrowingContinuation { (continuation: CheckedContinuation) -> Void in 19 | performRequest(with: image, level: fast ? .fast : .accurate) { request, error in 20 | if let error = error { 21 | continuation.resume(throwing: error) 22 | } else { 23 | guard let observations = request.results as? [VNRecognizedTextObservation] else { 24 | continuation.resume(throwing: Failure("No results")) 25 | return 26 | } 27 | 28 | var transcript: String = "" 29 | for observation in observations { 30 | transcript.append(observation.topCandidates(1)[0].string) 31 | transcript.append("\n") 32 | } 33 | 34 | continuation.resume(with: .success(transcript)) 35 | } 36 | } 37 | } 38 | } 39 | 40 | func performRequest(with image: CGImage, level: VNRequestTextRecognitionLevel, completion: @escaping VNRequestCompletionHandler) { 41 | let newHandler = VNImageRequestHandler(cgImage: image) 42 | 43 | let newRequest = VNRecognizeTextRequest(completionHandler: completion) 44 | newRequest.recognitionLevel = level 45 | 46 | do { 47 | if let customLanguages = try resolveLanguages(for: newRequest) { 48 | newRequest.recognitionLanguages = customLanguages 49 | } 50 | } catch { 51 | completion(newRequest, error) 52 | return 53 | } 54 | 55 | request = newRequest 56 | handler = newHandler 57 | 58 | do { 59 | try newHandler.perform([newRequest]) 60 | } catch { 61 | completion(newRequest, error) 62 | } 63 | } 64 | 65 | private func resolveLanguages(for request: VNRecognizeTextRequest) throws -> [String]? { 66 | try request.validateLanguages(with: customLanguages) 67 | } 68 | } 69 | -------------------------------------------------------------------------------- /Sources/ocrit/Implementation/AppleTranslateOperation.swift: -------------------------------------------------------------------------------- 1 | import Foundation 2 | private import Translation 3 | private import SwiftUI 4 | 5 | @available(macOS 15.0, *) 6 | struct AppleTranslateOperation: TranslationOperation { 7 | let text: String 8 | let inputLanguage: String 9 | let outputLanguage: String 10 | 11 | func run() async throws -> TranslationResult { 12 | let translator = Translator(sourceLanguage: inputLanguage, targetLanguage: outputLanguage) 13 | 14 | let response = try await translator.run(text) 15 | 16 | return TranslationResult( 17 | sourceText: text, 18 | translatedText: response.targetText, 19 | inputLanguage: inputLanguage, 20 | outputLanguage: outputLanguage 21 | ) 22 | } 23 | 24 | static func availability(from inputLanguage: String, to outputLanguage: String) async -> TranslationAvailability { 25 | let availability = LanguageAvailability() 26 | let status = await availability.status(from: .init(identifier: inputLanguage), to: .init(identifier: outputLanguage)) 27 | 28 | return switch status { 29 | case .supported: .supported 30 | case .installed: .installed 31 | case .unsupported: .unsupported 32 | @unknown default: .unsupported 33 | } 34 | } 35 | } 36 | 37 | // MARK: - Translation Shenanigans 38 | 39 | /** 40 | So, here's the thing: Translation was REALLY not meant to be run outside of an app's user interface, 41 | but I also REALLY wanted this capability in OCRIT, so I did what I had to do. Don't judge me. 42 | */ 43 | @available(macOS 15.0, *) 44 | @MainActor 45 | private struct Translator { 46 | let sourceLanguage: String 47 | let targetLanguage: String 48 | 49 | private struct _UIShim: View { 50 | var sourceLanguage: String 51 | var targetLanguage: String 52 | var text: String 53 | var callback: (Result) -> () 54 | 55 | var body: some View { 56 | EmptyView() 57 | .translationTask(source: .init(identifier: sourceLanguage), target: .init(identifier: targetLanguage)) { session in 58 | do { 59 | let result = try await session.translate(text) 60 | callback(.success(result)) 61 | } catch { 62 | callback(.failure(error)) 63 | } 64 | } 65 | } 66 | } 67 | 68 | func run(_ text: String) async throws -> TranslationSession.Response { 69 | try await withCheckedThrowingContinuation { continuation in 70 | let shim = _UIShim(sourceLanguage: sourceLanguage, targetLanguage: targetLanguage, text: text) { 71 | continuation.resume(with: $0) 72 | } 73 | 74 | /// This somehow works when running from a SPM-based executable... 75 | let window = NSWindow(contentViewController: NSHostingController(rootView: shim)) 76 | window.setFrame(.zero, display: false) 77 | window.alphaValue = 0 78 | window.makeKeyAndOrderFront(nil) 79 | } 80 | } 81 | } 82 | 83 | @available(macOS 15.0, *) 84 | extension TranslationSession: @retroactive @unchecked Sendable { } 85 | -------------------------------------------------------------------------------- /Sources/ocrit/OCRIT.swift: -------------------------------------------------------------------------------- 1 | import Foundation 2 | import ArgumentParser 3 | import UniformTypeIdentifiers 4 | import class Vision.VNRecognizeTextRequest 5 | import PathKit 6 | 7 | struct Failure: LocalizedError, CustomStringConvertible { 8 | var errorDescription: String? 9 | init(_ desc: String) { self.errorDescription = desc } 10 | var description: String { errorDescription ?? "" } 11 | } 12 | 13 | @main 14 | struct ocrit: AsyncParsableCommand { 15 | 16 | @Argument(help: "Path or list of paths for the images") 17 | var imagePaths: [Path] 18 | 19 | @Option( 20 | name: .shortAndLong, help: "Path to a directory where the txt files will be written to, or - for standard output" 21 | ) 22 | var output: Output = .stdOutput 23 | 24 | @Option(name: .shortAndLong, help: "Language code to use for the recognition, can be repeated to select multiple languages") 25 | var language: [String] = [] 26 | 27 | @Option(name: [.customShort("t"), .customLong("translate")], help: """ 28 | Language code to translate the detected text into. Requires macOS 15 or later. 29 | 30 | When using this option, the source language must be specified with -l/--language. 31 | 32 | ⚠️ This feature is experimental, use at your own risk. 33 | """) 34 | var translateIntoLanguageCode: String? 35 | 36 | @Flag(name: .shortAndLong, help: """ 37 | When -t/--translate is used, delete the original text files and only keep the translated ones. 38 | 39 | Also omits printing original untranslated text when output is stdout. 40 | """) 41 | var deleteOriginals = false 42 | 43 | @Flag(name: .shortAndLong, help: "Uses an OCR algorithm that prioritizes speed over accuracy") 44 | var fast = false 45 | 46 | func validate() throws { 47 | if let path = output.path, !path.isDirectory { 48 | do { 49 | try path.mkdir() 50 | } catch { 51 | throw ValidationError("Output path doesn't exist or is not a directory, and a directory couldn't be created at \(output). \(error)") 52 | } 53 | } 54 | 55 | /// Validate languages before attempting any OCR operations so that we can exit early in case there's an unsupported language. 56 | try VNRecognizeTextRequest.validateLanguages(with: language) 57 | 58 | guard translateIntoLanguageCode != nil else { return } 59 | 60 | guard #available(macOS 15.0, *) else { 61 | throw ValidationError("Translation is only available in macOS 15 or later.") 62 | } 63 | guard !language.isEmpty else { 64 | throw ValidationError("When using -t/--translate, the language of the document must be specified with -l/--language.") 65 | } 66 | guard language.count == 1 else { 67 | throw ValidationError("When using -t/--translate, only a single language can be specified with -l/--language.") 68 | } 69 | } 70 | 71 | private func checkTranslationAvailability() async throws { 72 | guard #available(macOS 15.0, *) else { return } 73 | guard let translateIntoLanguageCode else { return } 74 | 75 | let availability = await AppleTranslateOperation.availability(from: language[0], to: translateIntoLanguageCode) 76 | 77 | switch availability { 78 | case .supported: 79 | throw ValidationError("Translation is not supported from \"\(language[0])\" to \"\(translateIntoLanguageCode)\".") 80 | case .installed: 81 | break 82 | case .unsupported: 83 | throw ValidationError(""" 84 | In order to translate from \"\(language[0])\" to \"\(translateIntoLanguageCode)\", language support must be installed on your system. 85 | 86 | Go to System Settings > Language & Region > Translation Languages to install the languages. 87 | """) 88 | } 89 | } 90 | 91 | func run() async throws { 92 | try await checkTranslationAvailability() 93 | 94 | let imageURLs = imagePaths.map(\.url) 95 | 96 | fputs("Validating images…\n", stderr) 97 | 98 | var operationType: OCROperation.Type = ImageOCROperation.self 99 | 100 | do { 101 | for url in imageURLs { 102 | guard FileManager.default.fileExists(atPath: url.path) else { 103 | throw Failure("Image doesn't exist at \(url.path)") 104 | } 105 | 106 | guard let type = (try url.resourceValues(forKeys: [.contentTypeKey])).contentType else { 107 | throw Failure("Unable to determine file type at \(url.path)") 108 | } 109 | 110 | if type.conforms(to: .image) { 111 | operationType = ImageOCROperation.self 112 | } else if type.conforms(to: .pdf) { 113 | operationType = PDFOCROperation.self 114 | } else { 115 | throw Failure("File type at \(url.path) is not supported: \(type.identifier)") 116 | } 117 | } 118 | } catch { 119 | fputs("WARN: \(error.localizedDescription)\n", stderr) 120 | } 121 | 122 | if language.isEmpty { 123 | fputs("Performing OCR…\n", stderr) 124 | } else { 125 | if language.count == 1 { 126 | fputs("Performing OCR with language: \(language[0])…\n", stderr) 127 | } else { 128 | fputs("Performing OCR with languages: \(language.joined(separator: ", "))…\n", stderr) 129 | } 130 | } 131 | 132 | for url in imageURLs { 133 | let operation = operationType.init(fileURL: url, customLanguages: language) 134 | let fileAttributes = try? url.resourceValues(forKeys: [.creationDateKey, .contentModificationDateKey]) 135 | 136 | do { 137 | for try await result in try operation.run(fast: fast) { 138 | let translation = await runTranslationIfNeeded(for: result) 139 | 140 | try await processResult( 141 | result, 142 | translation: translation, 143 | for: url, 144 | fileAttributes: fileAttributes 145 | ) 146 | } 147 | } catch { 148 | /// Exit with error if there's only one image, otherwise we won't interrupt execution and will keep trying the other ones. 149 | guard imageURLs.count > 1 else { 150 | throw error 151 | } 152 | 153 | fputs("OCR failed for \(url.lastPathComponent): \(error.localizedDescription)\n", stderr) 154 | } 155 | } 156 | } 157 | 158 | private func runTranslationIfNeeded(for result: OCRResult) async -> TranslationResult? { 159 | guard #available(macOS 15.0, *) else { return nil } 160 | guard let translateIntoLanguageCode else { return nil } 161 | 162 | do { 163 | let operation = AppleTranslateOperation( 164 | text: result.text, 165 | inputLanguage: language[0], 166 | outputLanguage: translateIntoLanguageCode 167 | ) 168 | 169 | let result = try await operation.run() 170 | 171 | return result 172 | } catch { 173 | fputs("WARN: Translation failed for \"\(result.suggestedFilename)\". \(error)", stderr) 174 | return nil 175 | } 176 | } 177 | 178 | private func processResult(_ result: OCRResult, translation: TranslationResult?, for imageURL: URL, fileAttributes: URLResourceValues?) async throws { 179 | guard let outputDirectoryURL = output.path?.url else { 180 | writeStandardOutput(for: result, translation: translation, imageURL: imageURL) 181 | return 182 | } 183 | 184 | let outputFileURL = outputDirectoryURL 185 | .appendingPathComponent(result.suggestedFilename) 186 | .appendingPathExtension("txt") 187 | 188 | try result.text.write(to: outputFileURL, atomically: true, encoding: .utf8) 189 | 190 | outputFileURL.mergeAttributes(fileAttributes) 191 | 192 | guard let translation else { return } 193 | 194 | let translatedFilename = outputFileURL 195 | .deletingPathExtension() 196 | .lastPathComponent + "_\(translation.outputLanguage)" 197 | 198 | let translatedURL = outputFileURL 199 | .deletingLastPathComponent() 200 | .appendingPathComponent(translatedFilename, conformingTo: .plainText) 201 | 202 | try translation.translatedText.write( 203 | to: translatedURL, 204 | atomically: true, 205 | encoding: .utf8 206 | ) 207 | 208 | translatedURL.mergeAttributes(fileAttributes) 209 | 210 | outputFileURL.deleteIfNeeded(deleteOriginals) 211 | } 212 | 213 | private func writeStandardOutput(for result: OCRResult, translation: TranslationResult?, imageURL: URL) { 214 | if let translation { 215 | /// Don't print out original untranslated OCR if -d is specified. 216 | if !deleteOriginals { 217 | print("\(imageURL.lastPathComponent) (\(translation.inputLanguage))" + ":") 218 | print(result.text.trimmingCharacters(in: .whitespacesAndNewlines) + "\n") 219 | } 220 | 221 | print("\(imageURL.lastPathComponent) (\(translation.outputLanguage))" + ":") 222 | print(translation.translatedText.trimmingCharacters(in: .whitespacesAndNewlines) + "\n") 223 | } else { 224 | print(imageURL.lastPathComponent + ":") 225 | print(result.text.trimmingCharacters(in: .whitespacesAndNewlines) + "\n") 226 | } 227 | } 228 | } 229 | 230 | // MARK: - Helpers 231 | 232 | private extension URL { 233 | /// We don't want the entire OCR operation to fail if the tool can't copy file attributes from the image into the OCRed txt. 234 | /// This function attempts to merge the attributes and just logs to stderr if that fails. 235 | func mergeAttributes(_ values: URLResourceValues?) { 236 | guard let values else { return } 237 | 238 | var mSelf = self 239 | 240 | do { 241 | try mSelf.setResourceValues(values) 242 | } catch { 243 | fputs("WARN: Failed to set file attributes for \"\(lastPathComponent)\". \(error)\n", stderr) 244 | } 245 | } 246 | 247 | /// Syntactic sugar for conditionally deleting a file with a non-fatal error. 248 | func deleteIfNeeded(_ delete: Bool) { 249 | guard delete else { return } 250 | 251 | do { 252 | try FileManager.default.removeItem(at: self) 253 | } catch { 254 | fputs("WARN: Failed to delete \"\(lastPathComponent)\". \(error)\n", stderr) 255 | } 256 | } 257 | } 258 | -------------------------------------------------------------------------------- /Tests/ocritTests/ocritTests.swift: -------------------------------------------------------------------------------- 1 | import XCTest 2 | import class Foundation.Bundle 3 | import ArgumentParser 4 | 5 | final class OCRITTests: XCTestCase { 6 | 7 | func testOutputToStdout() throws { 8 | let input = fixturePath(named: "test-en.png") 9 | 10 | try AssertExecuteCommand( 11 | command: "ocrit \(input)", 12 | expected: .stdout(.equal("test-en.png:\nSome text in English")), 13 | exitCode: .success 14 | ) 15 | } 16 | 17 | func testLanguageSelectionEnglish() throws { 18 | let input = fixturePath(named: "test-en.png") 19 | 20 | try AssertExecuteCommand( 21 | command: "ocrit \(input) --language en-US", 22 | expected: [ 23 | .stderr(.contain("en-US")), /// should print selected language to stderr 24 | .stdout(.contain("Some text in English")) /// should print correct OCR result to stdout 25 | ], 26 | exitCode: .success 27 | ) 28 | } 29 | 30 | func testLanguageSelectionPortuguese() throws { 31 | let input = fixturePath(named: "test-pt.png") 32 | 33 | try AssertExecuteCommand( 34 | command: "ocrit \(input) --language pt-BR", 35 | expected: [ 36 | .stderr(.contain("pt-BR")), /// should print selected language to stderr 37 | .stdout(.contain("Um texto em Português")) /// should print correct OCR result to stdout 38 | ], 39 | exitCode: .success 40 | ) 41 | } 42 | 43 | func testLanguageSelectionChinese() throws { 44 | let input = fixturePath(named: "test-zh.png") 45 | 46 | try AssertExecuteCommand( 47 | command: "ocrit \(input) --language zh-Hans", 48 | expected: [ 49 | .stderr(.contain("zh-Hans")), /// should print selected language to stderr 50 | .stdout(.contain("一些中文文本")) /// should print correct OCR result to stdout 51 | ], 52 | exitCode: .success 53 | ) 54 | } 55 | 56 | func testMultipleLanguages() throws { 57 | let input = fixturePath(named: "test-multi-en-ko.png") 58 | 59 | try AssertExecuteCommand( 60 | command: "ocrit \(input) -l ko-KR -l en-US", 61 | expected: [ 62 | .stdout(.contain("like turtles")), 63 | .stdout(.contain("나는 거북이를 좋아한다")), 64 | ], 65 | exitCode: .success 66 | ) 67 | } 68 | 69 | func testInvalidLanguageExitsWithError() throws { 70 | let input = fixturePath(named: "test-en.png") 71 | 72 | try AssertExecuteCommand( 73 | command: "ocrit \(input) --language someinvalidlanguage", 74 | expected: .stderr(.contain("Unsupported language")), 75 | exitCode: .validationFailure 76 | ) 77 | } 78 | 79 | func testSpeedOverAccuracy() throws { 80 | let input = fixturePath(named: "test-en.png") 81 | 82 | try AssertExecuteCommand( 83 | command: "ocrit \(input) --fast", 84 | expected: .stdout(.equal("test-en.png:\nSome text in English")), 85 | exitCode: .success 86 | ) 87 | } 88 | 89 | func testMultipleImagesOutputToDirectory() throws { 90 | let outputURL = try getScratchDirectory() 91 | 92 | print("[+] Scratch directory for this test case is at \(outputURL.path)") 93 | 94 | let expectations = [ 95 | ("test-en.png", "en-US", "test-en.txt", "Some text in English"), 96 | ("test-pt.png", "pt-BR", "test-pt.txt", "Um texto em Português"), 97 | ("test-zh.png", "zh-Hans", "test-zh.txt", "一些中文文本") 98 | ] 99 | 100 | for (inputFilename, language, outputFilename, text) in expectations { 101 | let input = fixturePath(named: inputFilename) 102 | 103 | try AssertExecuteCommand( 104 | command: "ocrit \(input) --language \(language) --output \(outputURL.path)", 105 | exitCode: .success 106 | ) 107 | 108 | let outURL = outputURL.appendingPathComponent(outputFilename) 109 | 110 | XCTAssertTrue(FileManager.default.fileExists(atPath: outURL.path), "Output file for \(language) wasn't written") 111 | 112 | let output = try String(contentsOf: outURL, encoding: .utf8) 113 | 114 | XCTAssertTrue(output.contains(text), "Output file for \(language) doesn't contain \(text): \(outURL.path)") 115 | } 116 | } 117 | 118 | func testSinglePagePDFOutputToStdout() throws { 119 | let input = fixturePath(named: "test-en-singlepage.pdf") 120 | 121 | try AssertExecuteCommand( 122 | command: "ocrit \(input)", 123 | expected: .stdout(.contain("You can update your iPhone to iOS 17.4.1 by heading to the Settings app")), 124 | exitCode: .success 125 | ) 126 | } 127 | 128 | func testMultipagePDFOutputToStdout() throws { 129 | let input = fixturePath(named: "test-en-multipage.pdf") 130 | 131 | try AssertExecuteCommand( 132 | command: "ocrit \(input)", 133 | expected: [ 134 | .stdout(.contain("You can update your iPhone to iOS 17.4.1 by heading to the Settings app")), /// From page 1 135 | .stdout(.contain("When you add a resource to your Swift package, Xcode detects common resource types")), /// From page 2 136 | .stdout(.contain("To add a resource that Xcode can't handle automatically")), /// From page 3 137 | ], 138 | exitCode: .success 139 | ) 140 | } 141 | 142 | func testMultipagePDFOutputToDirectory() throws { 143 | let outputURL = try getScratchDirectory() 144 | 145 | print("[+] Scratch directory for this test case is at \(outputURL.path)") 146 | 147 | let expectations = [ 148 | ("test-en-multipage-1.txt", "You can update your iPhone to iOS 17.4.1 by heading to the Settings app"), 149 | ("test-en-multipage-2.txt", "When you add a resource to your Swift package, Xcode detects common resource types"), 150 | ("test-en-multipage-3.txt", "To add a resource that Xcode can't handle automatically"), 151 | ] 152 | 153 | let input = fixturePath(named: "test-en-multipage.pdf") 154 | 155 | try AssertExecuteCommand( 156 | command: "ocrit \(input) --output \(outputURL.path)", 157 | exitCode: .success 158 | ) 159 | 160 | for (outputFilename, text) in expectations { 161 | let outURL = outputURL.appendingPathComponent(outputFilename) 162 | 163 | XCTAssertTrue(FileManager.default.fileExists(atPath: outURL.path), "Output file \(outputFilename) wasn't written") 164 | 165 | let output = try String(contentsOf: outURL, encoding: .utf8) 166 | 167 | XCTAssertTrue(output.localizedCaseInsensitiveContains(text), "Output file \(outputFilename) doesn't contain \(text): \(outURL.path)") 168 | } 169 | } 170 | 171 | func testBasicTranslation() throws { 172 | guard #available(macOS 15.0, *) else { 173 | fputs("\(#function) skipped because this feature requires macOS 15 or later.\n", stderr) 174 | return 175 | } 176 | 177 | flakyWarn() 178 | 179 | let input = fixturePath(named: "test-en.png") 180 | 181 | try AssertExecuteCommand( 182 | command: "ocrit \(input) -l en-US -t pt-BR", 183 | expected: .stdout(.equal(""" 184 | test-en.png (en-US): 185 | Some text in English 186 | 187 | test-en.png (pt-BR): 188 | Algum texto em inglês 189 | """)), 190 | exitCode: .success 191 | ) 192 | } 193 | 194 | func testBasicTranslationChinese() throws { 195 | guard #available(macOS 15.0, *) else { 196 | fputs("\(#function) skipped because this feature requires macOS 15 or later.\n", stderr) 197 | return 198 | } 199 | 200 | flakyWarn() 201 | 202 | let input = fixturePath(named: "test-zh.png") 203 | 204 | try AssertExecuteCommand( 205 | command: "ocrit \(input) --language zh-Hans --translate en-US", 206 | expected: [ 207 | .stderr(.contain("zh-Hans")), /// should print selected language to stderr 208 | .stdout(.contain("一些中文文本")), /// should print correct OCR result to stdout 209 | .stdout(.contain("Some Chinese texts")) /// should print translated result to stdout 210 | ], 211 | exitCode: .success 212 | ) 213 | } 214 | 215 | func testBasicTranslationDeleteOriginals() throws { 216 | guard #available(macOS 15.0, *) else { 217 | fputs("\(#function) skipped because this feature requires macOS 15 or later.\n", stderr) 218 | return 219 | } 220 | 221 | flakyWarn() 222 | 223 | let input = fixturePath(named: "test-en.png") 224 | 225 | /// When -d is specified, only the translated text should be written to stdout. 226 | try AssertExecuteCommand( 227 | command: "ocrit \(input) -l en-US -t pt-BR -d", 228 | expected: .stdout(.equal(""" 229 | test-en.png (pt-BR): 230 | Algum texto em inglês 231 | """)), 232 | exitCode: .success 233 | ) 234 | } 235 | 236 | func testPDFTranslation() throws { 237 | guard #available(macOS 15.0, *) else { 238 | fputs("\(#function) skipped because this feature requires macOS 15 or later.\n", stderr) 239 | return 240 | } 241 | 242 | flakyWarn() 243 | 244 | let outputURL = try getScratchDirectory() 245 | 246 | print("[+] Scratch directory for this test case is at \(outputURL.path)") 247 | 248 | let expectations = [ 249 | ("test-en-multipage-1.txt", "You can update your iPhone to iOS 17.4.1 by heading to the Settings app"), 250 | ("test-en-multipage-2.txt", "When you add a resource to your Swift package, Xcode detects common resource types"), 251 | ("test-en-multipage-3.txt", "To add a resource that Xcode can't handle automatically"), 252 | 253 | ("test-en-multipage-1_pt-BR.txt", "Você pode atualizar seu iPhone para o iOS 17.4.1 indo para o aplicativo Configurações"), 254 | ("test-en-multipage-2_pt-BR.txt", "Quando você adiciona um recurso ao seu pacote Swift"), 255 | ("test-en-multipage-3_pt-BR.txt", "Para adicionar um recurso que o Xcode não pode manipular automaticamente"), 256 | ] 257 | 258 | let input = fixturePath(named: "test-en-multipage.pdf") 259 | 260 | try AssertExecuteCommand( 261 | command: "ocrit \(input) --output \(outputURL.path) -l en-US -t pt-BR", 262 | exitCode: .success 263 | ) 264 | 265 | for (outputFilename, text) in expectations { 266 | let outURL = outputURL.appendingPathComponent(outputFilename) 267 | 268 | XCTAssertTrue(FileManager.default.fileExists(atPath: outURL.path), "Output file \(outputFilename) wasn't written") 269 | 270 | let output = try String(contentsOf: outURL, encoding: .utf8) 271 | 272 | XCTAssertTrue(output.localizedCaseInsensitiveContains(text), "Output file \(outputFilename) doesn't contain \(text): \(outURL.path)") 273 | } 274 | 275 | } 276 | 277 | private func flakyWarn() { 278 | fputs("WARNING: Translation tests are likely going to be flaky as translation models get updated by Apple.\n", stderr) 279 | } 280 | } 281 | 282 | private extension OCRITTests { 283 | func fixturePath(named name: String) -> String { 284 | guard let url = Bundle.module.url(forResource: name, withExtension: "") else { 285 | XCTFail("Couldn't locate fixture \(name)") 286 | fatalError() 287 | } 288 | 289 | return url.path 290 | } 291 | 292 | func getScratchDirectory() throws -> URL { 293 | let outputDir = URL(fileURLWithPath: NSTemporaryDirectory()) 294 | .appendingPathComponent("ocrit_tests_\(UUID())") 295 | 296 | try FileManager.default.createDirectory(at: outputDir, withIntermediateDirectories: true) 297 | 298 | return outputDir 299 | } 300 | 301 | } 302 | 303 | extension String { 304 | var quoted: String { "\"\(self)\"" } 305 | } 306 | -------------------------------------------------------------------------------- /Tests/ocritTests/TestHelpers.swift: -------------------------------------------------------------------------------- 1 | //===----------------------------------------------------------*- swift -*-===// 2 | // 3 | // This source file is part of the Swift Argument Parser open source project 4 | // 5 | // Copyright (c) 2020 Apple Inc. and the Swift project authors 6 | // Licensed under Apache License v2.0 with Runtime Library Exception 7 | // 8 | // See https://swift.org/LICENSE.txt for license information 9 | // 10 | //===----------------------------------------------------------------------===// 11 | 12 | @testable import ArgumentParser 13 | import ArgumentParserToolInfo 14 | import XCTest 15 | 16 | // extensions to the ParsableArguments protocol to facilitate XCTestExpectation support 17 | public protocol TestableParsableArguments: ParsableArguments { 18 | var didValidateExpectation: XCTestExpectation { get } 19 | } 20 | 21 | public extension TestableParsableArguments { 22 | mutating func validate() throws { 23 | didValidateExpectation.fulfill() 24 | } 25 | } 26 | 27 | // extensions to the ParsableCommand protocol to facilitate XCTestExpectation support 28 | public protocol TestableParsableCommand: ParsableCommand, TestableParsableArguments { 29 | var didRunExpectation: XCTestExpectation { get } 30 | } 31 | 32 | public extension TestableParsableCommand { 33 | mutating func run() throws { 34 | didRunExpectation.fulfill() 35 | } 36 | } 37 | 38 | extension XCTestExpectation { 39 | public convenience init(singleExpectation description: String) { 40 | self.init(description: description) 41 | expectedFulfillmentCount = 1 42 | assertForOverFulfill = true 43 | } 44 | } 45 | 46 | public func AssertResultFailure( 47 | _ expression: @autoclosure () -> Result, 48 | _ message: @autoclosure () -> String = "", 49 | file: StaticString = #file, 50 | line: UInt = #line) 51 | { 52 | switch expression() { 53 | case .success: 54 | let msg = message() 55 | XCTFail(msg.isEmpty ? "Incorrectly succeeded" : msg, file: (file), line: line) 56 | case .failure: 57 | break 58 | } 59 | } 60 | 61 | public func AssertErrorMessage(_ type: A.Type, _ arguments: [String], _ errorMessage: String, file: StaticString = #file, line: UInt = #line) where A: ParsableArguments { 62 | do { 63 | _ = try A.parse(arguments) 64 | XCTFail("Parsing should have failed.", file: (file), line: line) 65 | } catch { 66 | // We expect to hit this path, i.e. getting an error: 67 | XCTAssertEqual(A.message(for: error), errorMessage, file: (file), line: line) 68 | } 69 | } 70 | 71 | public func AssertFullErrorMessage(_ type: A.Type, _ arguments: [String], _ errorMessage: String, file: StaticString = #file, line: UInt = #line) where A: ParsableArguments { 72 | do { 73 | _ = try A.parse(arguments) 74 | XCTFail("Parsing should have failed.", file: (file), line: line) 75 | } catch { 76 | // We expect to hit this path, i.e. getting an error: 77 | XCTAssertEqual(A.fullMessage(for: error), errorMessage, file: (file), line: line) 78 | } 79 | } 80 | 81 | public func AssertParse(_ type: A.Type, _ arguments: [String], file: StaticString = #file, line: UInt = #line, closure: (A) throws -> Void) where A: ParsableArguments { 82 | do { 83 | let parsed = try type.parse(arguments) 84 | try closure(parsed) 85 | } catch { 86 | let message = type.message(for: error) 87 | XCTFail("\"\(message)\" — \(error)", file: (file), line: line) 88 | } 89 | } 90 | 91 | public func AssertParseCommand(_ rootCommand: ParsableCommand.Type, _ type: A.Type, _ arguments: [String], file: StaticString = #file, line: UInt = #line, closure: (A) throws -> Void) { 92 | do { 93 | let command = try rootCommand.parseAsRoot(arguments) 94 | guard let aCommand = command as? A else { 95 | XCTFail("Command is of unexpected type: \(command)", file: (file), line: line) 96 | return 97 | } 98 | try closure(aCommand) 99 | } catch { 100 | let message = rootCommand.message(for: error) 101 | XCTFail("\"\(message)\" — \(error)", file: (file), line: line) 102 | } 103 | } 104 | 105 | public func AssertEqualStringsIgnoringTrailingWhitespace(_ string1: String, _ string2: String, file: StaticString = #file, line: UInt = #line) { 106 | let lines1 = string1.split(separator: "\n", omittingEmptySubsequences: false) 107 | let lines2 = string2.split(separator: "\n", omittingEmptySubsequences: false) 108 | 109 | XCTAssertEqual(lines1.count, lines2.count, "Strings have different numbers of lines.", file: (file), line: line) 110 | for (line1, line2) in zip(lines1, lines2) { 111 | XCTAssertEqual(line1.trimmed(), line2.trimmed(), file: (file), line: line) 112 | } 113 | } 114 | 115 | public func AssertHelp( 116 | _ visibility: ArgumentVisibility, 117 | for _: T.Type, 118 | equals expected: String, 119 | file: StaticString = #file, 120 | line: UInt = #line 121 | ) { 122 | let flag: String 123 | let includeHidden: Bool 124 | 125 | switch visibility.base { 126 | case .default: 127 | flag = "--help" 128 | includeHidden = false 129 | case .hidden: 130 | flag = "--help-hidden" 131 | includeHidden = true 132 | case .private: 133 | XCTFail("Should not be called.") 134 | return 135 | } 136 | 137 | do { 138 | _ = try T.parse([flag]) 139 | XCTFail(file: file, line: line) 140 | } catch { 141 | let helpString = T.fullMessage(for: error) 142 | AssertEqualStringsIgnoringTrailingWhitespace( 143 | helpString, expected, file: file, line: line) 144 | } 145 | 146 | let helpString = T.helpMessage(includeHidden: includeHidden, columns: nil) 147 | AssertEqualStringsIgnoringTrailingWhitespace( 148 | helpString, expected, file: file, line: line) 149 | } 150 | 151 | public func AssertHelp( 152 | _ visibility: ArgumentVisibility, 153 | for _: T.Type, 154 | root _: U.Type, 155 | equals expected: String, 156 | file: StaticString = #file, 157 | line: UInt = #line 158 | ) { 159 | let includeHidden: Bool 160 | 161 | switch visibility.base { 162 | case .default: 163 | includeHidden = false 164 | case .hidden: 165 | includeHidden = true 166 | case .private: 167 | XCTFail("Should not be called.") 168 | return 169 | } 170 | 171 | let helpString = U.helpMessage( 172 | for: T.self, includeHidden: includeHidden, columns: nil) 173 | AssertEqualStringsIgnoringTrailingWhitespace( 174 | helpString, expected, file: file, line: line) 175 | } 176 | 177 | public func AssertDump( 178 | for _: T.Type, equals expected: String, 179 | file: StaticString = #file, line: UInt = #line 180 | ) throws { 181 | do { 182 | _ = try T.parse(["--experimental-dump-help"]) 183 | XCTFail(file: (file), line: line) 184 | } catch { 185 | let dumpString = T.fullMessage(for: error) 186 | try AssertJSONEqualFromString(actual: dumpString, expected: expected, for: ToolInfoV0.self) 187 | } 188 | 189 | try AssertJSONEqualFromString(actual: T._dumpHelp(), expected: expected, for: ToolInfoV0.self) 190 | } 191 | 192 | public func AssertJSONEqualFromString(actual: String, expected: String, for type: T.Type) throws { 193 | let actualJSONData = try XCTUnwrap(actual.data(using: .utf8)) 194 | let actualDumpJSON = try XCTUnwrap(JSONDecoder().decode(type, from: actualJSONData)) 195 | 196 | let expectedJSONData = try XCTUnwrap(expected.data(using: .utf8)) 197 | let expectedDumpJSON = try XCTUnwrap(JSONDecoder().decode(type, from: expectedJSONData)) 198 | XCTAssertEqual(actualDumpJSON, expectedDumpJSON) 199 | } 200 | 201 | public enum ExpectedOutput: Hashable { 202 | public enum Operator: Hashable { 203 | case equal(String) 204 | case notEqual(String) 205 | case contain(String) 206 | case notContain(String) 207 | 208 | var expectedValue: String { 209 | switch self { 210 | case .equal(let v), .notEqual(let v), .contain(let v), .notContain(let v): 211 | return v 212 | } 213 | } 214 | } 215 | 216 | case stderr(Operator) 217 | case stdout(Operator) 218 | } 219 | 220 | extension XCTest { 221 | public var debugURL: URL { 222 | let bundleURL = Bundle(for: type(of: self)).bundleURL 223 | return bundleURL.lastPathComponent.hasSuffix("xctest") 224 | ? bundleURL.deletingLastPathComponent() 225 | : bundleURL 226 | } 227 | 228 | func validateOutput(stderr: String, stdout: String, expected: Set, file: StaticString, line: UInt) { 229 | func validate(value: String, op: ExpectedOutput.Operator) { 230 | let sanitizedValue = value.trimmingCharacters(in: .whitespacesAndNewlines) 231 | let expectedValue = op.expectedValue 232 | let sanitizedExpectedValue = expectedValue.trimmingCharacters(in: .whitespacesAndNewlines) 233 | 234 | switch op { 235 | case .equal: 236 | if sanitizedValue != sanitizedExpectedValue { 237 | XCTFail("\"\(value)\" is not equal to \"\(expectedValue)\"", 238 | file: (file), line: line) 239 | } 240 | case .notEqual: 241 | if sanitizedValue == sanitizedExpectedValue { 242 | XCTFail("\"\(value)\" is equal to \"\(expectedValue)\"", 243 | file: (file), line: line) 244 | } 245 | case .contain: 246 | if !sanitizedValue.localizedCaseInsensitiveContains(sanitizedExpectedValue) { 247 | XCTFail("\"\(value)\" does not contain \"\(expectedValue)\"", 248 | file: (file), line: line) 249 | } 250 | case .notContain: 251 | if sanitizedValue.localizedCaseInsensitiveContains(sanitizedExpectedValue) { 252 | XCTFail("\"\(value)\" contains \"\(expectedValue)\"", 253 | file: (file), line: line) 254 | } 255 | } 256 | } 257 | 258 | for expect in expected { 259 | switch expect { 260 | case .stderr(let op): 261 | validate(value: stderr, op: op) 262 | case .stdout(let op): 263 | validate(value: stdout, op: op) 264 | } 265 | } 266 | } 267 | 268 | public func AssertExecuteCommand( 269 | command: String, 270 | expected: ExpectedOutput, 271 | exitCode: ExitCode = .success, 272 | file: StaticString = #file, line: UInt = #line) throws 273 | { 274 | try AssertExecuteCommand(command: command, expected: [expected], exitCode: exitCode, file: file, line: line) 275 | } 276 | 277 | public func AssertExecuteCommand( 278 | command: String, 279 | expected: Set? = nil, 280 | exitCode: ExitCode = .success, 281 | file: StaticString = #file, line: UInt = #line) throws 282 | { 283 | let splitCommand = command.split(separator: " ") 284 | let arguments = splitCommand.dropFirst().map(String.init) 285 | 286 | let commandName = String(splitCommand.first!) 287 | let commandURL = debugURL.appendingPathComponent(commandName) 288 | guard (try? commandURL.checkResourceIsReachable()) ?? false else { 289 | XCTFail("No executable at '\(commandURL.standardizedFileURL.path)'.", 290 | file: (file), line: line) 291 | return 292 | } 293 | 294 | #if !canImport(Darwin) || os(macOS) 295 | let process = Process() 296 | if #available(macOS 10.13, *) { 297 | process.executableURL = commandURL 298 | } else { 299 | process.launchPath = commandURL.path 300 | } 301 | process.arguments = arguments 302 | 303 | let output = Pipe() 304 | process.standardOutput = output 305 | let error = Pipe() 306 | process.standardError = error 307 | 308 | if #available(macOS 10.13, *) { 309 | guard (try? process.run()) != nil else { 310 | XCTFail("Couldn't run command process.", file: (file), line: line) 311 | return 312 | } 313 | } else { 314 | process.launch() 315 | } 316 | process.waitUntilExit() 317 | 318 | let outputData = output.fileHandleForReading.readDataToEndOfFile() 319 | let outputActual = String(data: outputData, encoding: .utf8)!.trimmingCharacters(in: .whitespacesAndNewlines) 320 | 321 | let errorData = error.fileHandleForReading.readDataToEndOfFile() 322 | let errorActual = String(data: errorData, encoding: .utf8)!.trimmingCharacters(in: .whitespacesAndNewlines) 323 | 324 | if let expected = expected { 325 | validateOutput(stderr: errorActual, stdout: outputActual, expected: expected, file: file, line: line) 326 | } 327 | 328 | XCTAssertEqual(process.terminationStatus, exitCode.rawValue, file: (file), line: line) 329 | #else 330 | throw XCTSkip("Not supported on this platform") 331 | #endif 332 | } 333 | 334 | public func AssertJSONOutputEqual( 335 | command: String, 336 | expected: String, 337 | file: StaticString = #file, line: UInt = #line 338 | ) throws { 339 | 340 | let splitCommand = command.split(separator: " ") 341 | let arguments = splitCommand.dropFirst().map(String.init) 342 | 343 | let commandName = String(splitCommand.first!) 344 | let commandURL = debugURL.appendingPathComponent(commandName) 345 | guard (try? commandURL.checkResourceIsReachable()) ?? false else { 346 | XCTFail("No executable at '\(commandURL.standardizedFileURL.path)'.", 347 | file: (file), line: line) 348 | return 349 | } 350 | 351 | #if !canImport(Darwin) || os(macOS) 352 | let process = Process() 353 | if #available(macOS 10.13, *) { 354 | process.executableURL = commandURL 355 | } else { 356 | process.launchPath = commandURL.path 357 | } 358 | process.arguments = arguments 359 | 360 | let output = Pipe() 361 | process.standardOutput = output 362 | let error = Pipe() 363 | process.standardError = error 364 | 365 | if #available(macOS 10.13, *) { 366 | guard (try? process.run()) != nil else { 367 | XCTFail("Couldn't run command process.", file: (file), line: line) 368 | return 369 | } 370 | } else { 371 | process.launch() 372 | } 373 | process.waitUntilExit() 374 | 375 | let outputString = try XCTUnwrap(String(data: output.fileHandleForReading.readDataToEndOfFile(), encoding: .utf8)) 376 | XCTAssertTrue(error.fileHandleForReading.readDataToEndOfFile().isEmpty, "Error occurred with `--experimental-dump-help`") 377 | try AssertJSONEqualFromString(actual: outputString, expected: expected, for: ToolInfoV0.self) 378 | #else 379 | throw XCTSkip("Not supported on this platform") 380 | #endif 381 | } 382 | } 383 | --------------------------------------------------------------------------------