srtune/src/main.rs

#[macro_use]
extern crate log;
#[macro_use]
extern crate lazy_static;

use regex::Regex;
use std::fs::{File, OpenOptions};
use std::io::{BufReader, Write};
use std::io;
use std::str::FromStr;
use std::convert::TryFrom;
use std::io::BufRead;
use std::ops::{Add, Mul, MulAssign, AddAssign};
use std::fmt::{self, Display};

const LOG_LEVELS: [&str; 5] = ["error", "warn", "info", "debug", "trace"];
const SPAMMY_LIBS: [&str; 5] = ["tokio_reactor", "hyper", "reqwest", "mio", "want"];

fn main() {
    let argv =
        clap::App::new("srtune")
            .version(env!("CARGO_PKG_VERSION"))
            .arg(clap::Arg::with_name("input")
                     .value_name("INFILE")
                     .help("Input file, -- for stdin"),
            )
            .arg(clap::Arg::with_name("output")
                     .short("o")
                     .long("output")
                     .value_name("OUTFILE")
                     .help("Output file, defaults to stdout"),
            )
            .arg(clap::Arg::with_name("from")
                     .short("f")
                     .long("from")
                     .value_name("FROM")
                     .help("Index of the first affected entry, defaults to 0. With '--renumber', the original numbers are used."),
            )
            .arg(clap::Arg::with_name("move")
                     .short("m")
                     .long("move")
                     .value_name("SECS")
                     .help("Time shift, accepts positive or negative float"),
            )
            .arg(clap::Arg::with_name("scale")
                     .short("s")
                     .long("scale")
                     .value_name("RATIO")
                     .help("Scale subtitle times and durations to compensate for bitrate differences.\nIf a start time was given, the scaling is relative to this point, otherwise to the first subtitle in the file."),
            )
            .arg(clap::Arg::with_name("durscale")
                     .short("d")
                     .long("durscale")
                     .value_name("RATIO")
                     .help("Scale durations, can be combined with '--scale'"),
            )
            .arg(clap::Arg::with_name("renumber")
                     .short("r")
                     .long("renumber")
                     .help("Change all numbers to be sequential starting with 1"),
            )
            .arg(clap::Arg::with_name("v").short("v").multiple(true).help(
                "Sets the level of verbosity (adds to the default - info)",
            ))
            .get_matches();

    let mut log_level = "info".to_owned();

    if argv.is_present("v") {
        // bump verbosity if -v's are present
        let pos = LOG_LEVELS
            .iter()
            .position(|x| x == &log_level)
            .unwrap();

        log_level = match LOG_LEVELS
            .iter()
            .nth(pos + argv.occurrences_of("v") as usize)
            {
                Some(new_level) => new_level.to_string(),
                None => "trace".to_owned(),
            };
    }

    //println!("LEVEL={}", log_level);

    // init logging
    let env = env_logger::Env::default().default_filter_or(log_level);
    let mut builder = env_logger::Builder::from_env(env);
    let lib_level = log::LevelFilter::Info;
    for lib in &SPAMMY_LIBS {
        builder.filter_module(lib, lib_level);
    }
    builder.init();

    let inf = argv.value_of("input");
    let outf = argv.value_of("output");
    let stdin = io::stdin();
    let stdout = io::stdout();

    let mut iter: Box<dyn Iterator<Item=Result<String, io::Error>>> = match inf {
        None => {
            Box::new(stdin.lock().lines())
        }
        Some(f) => {
            let file = File::open(f).expect(&format!("Could not open file: {:?}", f));
            Box::new(BufReader::new(file).lines())
        }
    };

    let mut out : Box<dyn Write> = match outf {
        None => {
            Box::new(stdout.lock())
        }
        Some(f) => {
            let file = OpenOptions::new()
                .create(true)
                .truncate(true)
                .write(true)
                .open(f)
                .expect(&format!("Could not open file: {:?}", f));
            Box::new(file)
        }
    };

    let from_s = match argv.value_of("from") {
        Some(s) => {
            s.parse().expect("Bad --from format")
        },
        None => 0u32
    };

    let move_s = match argv.value_of("move") {
        Some(s) => {
            s.parse().expect("Bad --move format")
        },
        None => 0f32
    };

    let scale_s = match argv.value_of("scale") {
        Some(s) => {
            s.parse().expect("Bad --scale format")
        },
        None => 1f32
    };

    let durscale_s = match argv.value_of("durscale") {
        Some(s) => {
            s.parse().expect("Bad --durscale format")
        },
        None => 1f32
    } * scale_s; // always also shrink durations

    let renumber = argv.is_present("renumber");

    info!("Opts: from #{}, move {}s, scale {}x, durscale {}x", from_s, move_s, scale_s, durscale_s);

    let mut scale_start = SubTime(0f32);
    let mut first_found = false;
    let mut renumber_i : u32 = 0;

    let mut text = vec![];
    while let Some(Ok(x)) = iter.next() {
        let mut x = x.trim();
        if x.starts_with('\u{feff}') {
            debug!("Stripping BOM mark");
            x = &x[3..];
        }
        let x = x.trim();
        if x.is_empty() {
            continue;
        }

        // 236
        // 00:18:01,755 --> 00:18:03,774
        // （掃除機の音）
        // う～ん…。
        match u32::from_str(x) {
            Ok(num) => {
                // println!("Entry {}", num);
                let datesrow = iter.next().unwrap().unwrap();
                if datesrow.contains(" --> ") {
                    let mut halves = datesrow.split(" --> ");
                    let (first, second) = (halves.next().unwrap(), halves.next().unwrap());
                    let start = SubTime::try_from(first).unwrap();
                    let end = SubTime::try_from(second).unwrap();

                    text.clear();
                    while let Some(Ok(x)) = iter.next() {
                        if x.is_empty() {
                            break; // space between the entries
                        }
                        text.push(x);
                    }

                    let mut one = Subtitle {
                        num,
                        start,
                        dur: SubDuration(end.0 - start.0),
                        text: text.join("\n"),
                    };

                    if num >= from_s {
                        if !first_found {
                            debug!("Scaling anchored at {} (#{}), start edits", start, num);
                            scale_start = start;
                            first_found = true;
                        }

                        if scale_s != 1f32 {
                            one.start = one.start.scale(scale_start, scale_s);
                        }

                        one.dur *= durscale_s;
                        one.start += move_s;
                    }

                    if one.start.0 < 0f32 {
                        warn!("Discarding negative time entry #{} @ {:.3}s", one.num, one.start.0);
                        continue;
                    }

                    // advance numbering only for the really emitted entries
                    if renumber {
                        renumber_i += 1;
                        one.num = renumber_i;
                    }

                    out.write(one.to_string().as_bytes()).expect("failed to write");
                }
            }
            Err(e) => {
                error!("couldnt parse >{}<: {}", x, e);
                for b in x.as_bytes() {
                    error!("{:#02x} - {}", b, b);
                }
                error!("\n");
            }
        }
    }

    out.flush().unwrap();
}

#[derive(Copy, Clone, Debug)]
struct SubTime(f32);

#[derive(Copy, Clone, Debug)]
struct SubDuration(f32);

impl Add<SubDuration> for SubTime {
    type Output = SubTime;

    fn add(self, rhs: SubDuration) -> Self::Output {
        SubTime(self.0 + rhs.0)
    }
}

impl Mul<f32> for SubDuration {
    type Output = SubDuration;

    fn mul(self, rhs: f32) -> Self::Output {
        SubDuration(self.0 * rhs)
    }
}

impl MulAssign<f32> for SubDuration {
    fn mul_assign(&mut self, rhs: f32) {
        self.0 *= rhs;
    }
}

impl AddAssign<f32> for SubDuration {
    fn add_assign(&mut self, rhs: f32) {
        self.0 += rhs;
    }
}

impl SubTime {
    /// Scale by a factor with a custom start time
    pub fn scale(&self, start : SubTime, factor : f32) -> SubTime {
        SubTime(start.0 + (self.0 - start.0) * factor)
    }
}

impl AddAssign<f32> for SubTime {
    fn add_assign(&mut self, rhs: f32) {
        self.0 += rhs;
    }
}

impl Display for SubTime {
    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        // TODO optimize this
        let mut secs = self.0;
        let hours = (secs / 3600f32).floor();
        secs -= hours * 3600f32;
        let minutes = (secs / 60f32).floor();
        secs -= minutes * 60f32;
        let msecs = ((secs % 1f32)*1000f32).round();
        write!(f, "{:02}:{:02}:{:02},{:03}", hours, minutes, secs.floor(), msecs)
    }
}

#[derive(Clone, Debug)]
struct Subtitle {
    num: u32,
    start: SubTime,
    dur: SubDuration,
    text: String,
}

impl Display for Subtitle {
    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        write!(f, "{}\n{} --> {}\n{}\n\n",
               self.num,
               self.start, self.start + self.dur,
               self.text
        )
    }
}

lazy_static! {
    static ref DATE_RE: Regex = Regex::new(r"(\d+):(\d+):(\d+),(\d+)").unwrap();
}

impl TryFrom<&str> for SubTime {
    type Error = failure::Error;

    fn try_from(value: &str) -> Result<Self, Self::Error> {
        match DATE_RE.captures(value) {
            Some(caps) => {
                Ok(SubTime(f32::from_str(caps.get(1).unwrap().as_str()).unwrap() * 3600f32 +
                    f32::from_str(caps.get(2).unwrap().as_str()).unwrap() * 60f32 +
                    f32::from_str(caps.get(3).unwrap().as_str()).unwrap() +
                    f32::from_str(caps.get(4).unwrap().as_str()).unwrap() * 0.001f32))
            }
            None => Err(failure::err_msg("Error parsing time."))
        }
    }
}