Skip to main content

srt_writer/
lib.rs

1// Copyright (C) The Strand-Braid Authors
2// SPDX-License-Identifier: MIT OR Apache-2.0
3
4use std::{io::Result, io::Write, time::Duration};
5
6trait Srt {
7    fn srt(&self) -> String;
8}
9
10impl Srt for Duration {
11    fn srt(&self) -> String {
12        // from https://en.wikipedia.org/wiki/SubRip :
13        // "hours:minutes:seconds,milliseconds with time units fixed to two
14        // zero-padded digits and fractions fixed to three zero-padded digits
15        // (00:00:00,000). The fractional separator used is the comma, since the
16        // program was written in France."
17        let total_secs = self.as_secs();
18        let hours = total_secs / (60 * 60);
19        let minutes = (total_secs % (60 * 60)) / 60;
20        let seconds = total_secs % 60;
21        debug_assert_eq!(total_secs, hours * 60 * 60 + minutes * 60 + seconds);
22        let millis = self.subsec_millis();
23        format!("{hours:02}:{minutes:02}:{seconds:02},{millis:03}")
24    }
25}
26
27pub struct SrtWriter {
28    wtr: Box<dyn Write>,
29    count: usize,
30}
31
32impl SrtWriter {
33    pub fn new(wtr: Box<dyn Write>) -> Self {
34        Self { wtr, count: 1 }
35    }
36
37    pub fn append(&mut self, start: Duration, stop: Duration, value: &str) -> Result<()> {
38        self.wtr.write_all(
39            format!(
40                "{count}\n{start} --> {stop}\n{value}\n\n",
41                count = self.count,
42                start = start.srt(),
43                stop = stop.srt(),
44            )
45            .as_bytes(),
46        )?;
47        self.count += 1;
48        Ok(())
49    }
50
51    pub fn flush(&mut self) -> Result<()> {
52        self.wtr.flush()
53    }
54
55    pub fn close(mut self) -> Result<()> {
56        self.wtr.flush()
57    }
58}
59
60/// A buffering [SrtWriter] which is meant to be called for every frame.
61///
62/// This buffers values from each frame until the next frame. In this way, it
63/// can calculate start and stop times for each frame. The first call to
64/// [Self::add_frame] thus only stores the buffer, and the buffered value is
65/// written upon [Self::close] or [Self::drop].
66pub struct BufferingSrtFrameWriter {
67    srt_wtr: SrtWriter,
68    prev: Option<(Duration, String)>,
69}
70
71impl BufferingSrtFrameWriter {
72    pub fn new(wtr: Box<dyn Write>) -> Self {
73        Self {
74            srt_wtr: SrtWriter::new(wtr),
75            prev: None,
76        }
77    }
78    pub fn add_frame(&mut self, pts: Duration, val: String) -> Result<()> {
79        if let Some((prev_pts, prev_value)) = self.prev.take() {
80            // write buffered value
81            self.srt_wtr.append(prev_pts, pts, &prev_value)?;
82        }
83        // store current value
84        self.prev = Some((pts, val));
85        Ok(())
86    }
87
88    /// Flush the underlying writer.
89    ///
90    /// Note that this does not flush the currently buffered value, as that
91    /// would require creating a new timestamp.
92    pub fn flush(&mut self) -> Result<()> {
93        self.srt_wtr.flush()
94    }
95
96    pub fn close(mut self) -> Result<()> {
97        self.end_with_fake_timestamp()?;
98        // Ensure no further frames are appended by dropping self.
99        Ok(())
100    }
101
102    /// End the file (private method)
103    ///
104    /// As this adds a fake timestamp, we do want to drop self so that we do not
105    /// continue appending timestamps after the bad one. This fake timestamp
106    /// should be only the final timestamp and not in the middle of the file.
107    ///
108    /// The caller must ensure that no further frames are appended, e.g. by
109    /// dropping this instance of Self.
110    fn end_with_fake_timestamp(&mut self) -> Result<()> {
111        if let Some((pts, value)) = self.prev.take() {
112            // invent timestamp in the future
113            let future_pts = pts + Duration::from_secs(1);
114            self.srt_wtr.append(pts, future_pts, &value)?;
115        }
116        // As simply dropping self.srt_wtr will not flush it, we must manually do
117        // it.
118        self.srt_wtr.flush()?;
119        Ok(())
120    }
121}
122
123impl Drop for BufferingSrtFrameWriter {
124    fn drop(&mut self) {
125        self.end_with_fake_timestamp().unwrap();
126        // We ensure no further frames are appended because we are in drop()
127        // here.
128    }
129}