Fix cue parsing when using carriage return

mihai8804858 · May 21, 2024 · 484a438 · 484a438
1 parent 250e582
commit 484a438
Show file tree

Hide file tree

Showing 2 changed files with 58 additions and 9 deletions.
diff --git a/Sources/Extensions/String.swift b/Sources/Extensions/String.swift
@@ -1,16 +1,22 @@
+import Foundation
+
 extension String {
     func trimmingEdges(while predicate: (Element) throws -> Bool) rethrows -> String {
-        let trimmed = try trimmingPrefix(while: predicate).trimmingSuffix(while: predicate)
-        return String(trimmed)
+        try String(self[...].trimmingEdges(while: predicate))
     }
+}
 
-    func prefix(upTo substring: Substring) -> String {
-        guard let range = range(of: substring) else { return self }
-        return String(self[..<range.lowerBound])
+extension CharacterSet {
+    func containsUnicodeScalars(of character: Character) -> Bool {
+        character.unicodeScalars.allSatisfy(contains)
     }
 }
 
 extension Substring {
+    func trimmingEdges(while predicate: (Element) throws -> Bool) rethrows -> Substring {
+        try trimmingPrefix(while: predicate).trimmingSuffix(while: predicate)
+    }
+
     func trimmingSuffix(while predicate: (Element) throws -> Bool) rethrows -> Substring {
         var result = self
         while let last = result.last, try predicate(last) {
@@ -19,4 +25,27 @@ extension Substring {
 
         return result
     }
+
+    func prefix(upTo terminatorSet: CharacterSet, count: Int = 1, includeTerminator: Bool = false) -> Substring {
+        var result = Substring()
+        var terminator = Substring()
+        var occurences = 0
+        for char in self {
+            if terminatorSet.containsUnicodeScalars(of: char) {
+                terminator.append(char)
+                occurences += 1
+                if occurences == count { break }
+            } else {
+                result.append(contentsOf: terminator)
+                result.append(char)
+                terminator.removeAll()
+                occurences = 0
+            }
+        }
+        if includeTerminator {
+            result.append(contentsOf: terminator)
+        }
+
+        return result
+    }
 }
diff --git a/Sources/SRTParser.swift b/Sources/SRTParser.swift
@@ -36,7 +36,7 @@ struct CueParser: ParserPrinter {
             Whitespace(.horizontal)
             Whitespace(1..., .vertical)
             CueMetadataParser()
-            TextParser()
+            TextParser(upTo: .newlines, count: 2)
         }
     }
 }
@@ -50,7 +50,7 @@ struct CueMetadataParser: ParserPrinter {
                 CoordinatesParser()
             }
             Whitespace(.horizontal)
-            Whitespace(1, .vertical)
+            Whitespace(1..., .vertical)
             Optionally {
                 PositionParser()
             }
@@ -126,9 +126,29 @@ struct PositionParser: ParserPrinter {
 }
 
 struct TextParser: ParserPrinter {
+    let terminator: CharacterSet
+    let count: Int
+    let includeTerminator: Bool
+
+    init(upTo terminator: String, count: Int = 1, includeTerminator: Bool = false) {
+        self.terminator = CharacterSet(charactersIn: terminator)
+        self.count = count
+        self.includeTerminator = includeTerminator
+    }
+
+    init(upTo terminator: CharacterSet, count: Int = 1, includeTerminator: Bool = false) {
+        self.terminator = terminator
+        self.count = count
+        self.includeTerminator = includeTerminator
+    }
+
     func parse(_ input: inout Substring) throws -> SRT.StyledText {
-        let prefix = String(input).prefix(upTo: "\n\n")
-        let text = try StyledTextParser().parse(prefix)
+        let prefix = input.prefix(
+            upTo: terminator,
+            count: count,
+            includeTerminator: includeTerminator
+        )
+        let text = try StyledTextParser().parse(String(prefix))
         input.removeFirst(prefix.count)
 
         return text