From 6ceb30b2c6990a535ecc8c026ddf3fdc5083c800 Mon Sep 17 00:00:00 2001
From: Jordi Mas <jmas@softcatala.org>
Date: Fri, 20 Dec 2024 08:10:14 +0100
Subject: [PATCH] Fix: export correct subtitles when time is one digit

---
 open_dubbing/subtitles.py | 25 ++++++++++++++++---------
 tests/subtitles_test.py   | 35 ++++++++++++++++++++++++++++-------
 2 files changed, 44 insertions(+), 16 deletions(-)

diff --git a/open_dubbing/subtitles.py b/open_dubbing/subtitles.py
index f96ab47..f44fd7f 100644
--- a/open_dubbing/subtitles.py
+++ b/open_dubbing/subtitles.py
@@ -20,22 +20,29 @@
 class Subtitles:
 
     def write(self, *, utterance_metadata, directory, filename, translated):
-        srt_file_path = filename
-        srt_file_path = os.path.join(directory, srt_file_path)
+        srt_file_path = os.path.join(directory, filename)
 
         with open(srt_file_path, "w", encoding="utf-8") as subtitles_file:
             for i, utterance in enumerate(utterance_metadata):
-                start_time = str(timedelta(seconds=utterance["start"]))[:-3]
-                end_time = str(timedelta(seconds=utterance["end"]))[:-3]
-                start_time = start_time.replace(".", ",").zfill(12)
-                end_time = end_time.replace(".", ",").zfill(12)
+                start_time = self.format_srt_time(utterance["start"])
+                end_time = self.format_srt_time(utterance["end"])
+
                 idx = i + 1
                 srt_content = f"{idx}\n"
-                srt_content += (
-                    f"{start_time.replace('.', ',')} --> {end_time.replace('.', ',')}\n"
-                )
+                srt_content += f"{start_time} --> {end_time}\n"
 
                 text = utterance["translated_text"] if translated else utterance["text"]
                 srt_content += f"{text}\n\n"
                 subtitles_file.write(srt_content)
         return srt_file_path
+
+    @staticmethod
+    def format_srt_time(seconds):
+        """Format seconds as HH:MM:SS,mmm for SRT files."""
+        time_delta = timedelta(seconds=seconds)
+        total_seconds = int(time_delta.total_seconds())
+        milliseconds = int((time_delta.total_seconds() - total_seconds) * 1000)
+        hours = total_seconds // 3600
+        minutes = (total_seconds % 3600) // 60
+        seconds = total_seconds % 60
+        return f"{hours:02}:{minutes:02}:{seconds:02},{milliseconds:03}"
diff --git a/tests/subtitles_test.py b/tests/subtitles_test.py
index 81bb39e..3ae47eb 100644
--- a/tests/subtitles_test.py
+++ b/tests/subtitles_test.py
@@ -96,14 +96,35 @@ def test_write_dubbed(self):
             "",
         ]
 
-    def _get_srt(self):
+    def _get_utterances_small_digits(self):
         return [
+            {
+                "id": 1,
+                "start": 1,
+                "end": 3.1,
+                "text": "Good morning, my name is Jordi Mas.",
+                "translated_text": "Bon dia, el meu nom és Jordi Mas.",
+            },
+        ]
+
+    def test_write_original_small_digits(self):
+        subtitles = Subtitles()
+        srt_file = tempfile.NamedTemporaryFile(suffix=".srt", delete=False).name
+
+        directory = os.path.dirname(srt_file)
+        filename = os.path.basename(srt_file)
+        subtitles.write(
+            utterance_metadata=self._get_utterances_small_digits(),
+            directory=directory,
+            filename=filename,
+            translated=False,
+        )
+
+        lines = self._get_lines_from_file(srt_file)
+        os.remove(srt_file)
+        assert lines == [
             "1",
-            "00:00:01,262 --> 00:00:03,945",
-            "Bon dia, el meu nom és Jordi Mas.",
-            "",
-            "2",
-            "00:00:05,245 --> 00:00:06,629",
-            "Sóc de Barcelona.",
+            "00:00:01,000 --> 00:00:03,100",
+            "Good morning, my name is Jordi Mas.",
             "",
         ]