From fa2610e206af8896831078a25518e04211b86335 Mon Sep 17 00:00:00 2001
From: Ethan Dalool <git@voussoir.net>
Date: Sat, 24 Jun 2023 12:03:55 -0700
Subject: [PATCH] Use pysrt to provide better content search of srt files.

---
 search.py | 29 +++++++++++++++++++++++++++++
 1 file changed, 29 insertions(+)

diff --git a/search.py b/search.py
index 5f5d88c..385ad56 100644
--- a/search.py
+++ b/search.py
@@ -9,6 +9,10 @@ try:
     import winshell
 except ImportError:
     winshell = None
+try:
+    import pysrt
+except ImportError:
+    pysrt = None
 
 from voussoirkit import expressionmatch
 from voussoirkit import pathclass
@@ -90,6 +94,29 @@ def search_contents_generic(filepath, content_args):
     yield from results
     yield ''
 
+def _srt_format_line(line):
+    text = line.text.replace('\n', ' ')
+    timestamp = f'{line.start.hours:02d}:{line.start.minutes:02d}:{line.start.seconds:02d}:{line.start.milliseconds:03d}'
+    return f'{timestamp} {text}'
+
+def search_contents_srt(filepath, content_args):
+    try:
+        srtlines = pysrt.open(filepath.absolute_path)
+    except UnicodeDecodeError:
+        log.warn('%s experienced Unicode Error', filepath.absolute_path)
+        return
+
+    content_args['text'] = '\n'.join(_srt_format_line(line) for line in srtlines)
+
+    results = search(**content_args)
+    results = list(results)
+    if not results:
+        return
+
+    yield filepath.absolute_path
+    yield from results
+    yield ''
+
 def search_contents_windows_lnk(filepath, content_args):
     try:
         shortcut = winshell.Shortcut(filepath.absolute_path)
@@ -227,6 +254,8 @@ def search(
 
         if filepath.extension == 'lnk' and winshell:
             yield from search_contents_windows_lnk(filepath, content_args)
+        if filepath.extension == 'srt' and pysrt:
+            yield from search_contents_srt(filepath, content_args)
         else:
             yield from search_contents_generic(filepath, content_args)