monorepo/dezendorf/applications/offset-fixer/offset-fixer.py

#!/usr/bin/env python

import re
from datetime import timedelta
import argparse
import locale

import math

class SubtitleEntry:
  def __init__(self, entryNumber, text, startTime, endTime):
    # dd:dd:dd,ddd --> dd:dd:dd,ddd
    self.entryNumber = entryNumber
    self.text = text
    sH, sM, sS = startTime.replace(',', '.').split(':')
    self.startTime = int(sH) * 3600 + int(sM) * 60 + float(sS)
    sH, sM, sS = endTime.replace(',', '.').split(':')
    self.endTime = int(sH) * 3600 + int(sM) * 60 + float(sS)

  def formatEntry(self):
    entry = ""
    entry += "{}\n".format(self.entryNumber)
    entry += self.formatTimeString() + "\n"
    entry += "{}".format(self.text)
    return entry

  def formatTimeString(self):
    s = self.startTime
    startHours = s // 3600
    s = s - (startHours * 3600)
    startMinutes = s // 60
    s = s - (startMinutes * 60)
    startSeconds = s

    e = self.endTime
    endHours = e // 3600
    e = e - (endHours * 3600)
    endMinutes = e // 60
    e = e - (endMinutes * 60)
    endSeconds = e

    return '{:02.0f}:{:02.0f}:{:02.0f},{:03.0f} --> {:02.0f}:{:02.0f}:{:02.0f},{:03.0f}'.format(startHours, startMinutes, math.floor(startSeconds), (startSeconds % 1) * 1000 , endHours, endMinutes, math.floor(endSeconds), (endSeconds % 1) * 1000 )

  def shift(self, seconds):
    self.startTime = self.startTime + seconds    
    self.endTime = self.endTime + seconds

def main():
  Subtitles = []

  parser = argparse.ArgumentParser()

  parser.add_argument("--seconds", help="Decimal formatted seconds to adjust offsets by", type=float, required=True)
  parser.add_argument("--file", help="Subtitle file to shift", type=str, required=True)
  parser.add_argument("--inplace", required=False, help="Shift file in place, creating a backup", action='store_true')
  parser.add_argument("--outfile", required=False, help="New subtitle filename", type=str)
  parser.add_argument("--strip-first", required=False, help="Clear first subtitle entry", action='store_true')
  parser.add_argument("--strip-last", required=False, help="Clear last subtitle entry", action='store_true')

  args = parser.parse_args()

  if args.inplace is True:
      if args.outfile == None:
        args.outfile = str(args.file + ".backup")
      else:
        sys.exit(1)

  if args.outfile == None:
      args.outfile = str(args.file + ".new")

  print("Opening subtitle file: {}".format(args.file))

  with open(args.file, "r+") as file:
    entryNumber = 0
    text = ""
    timestamp = ""
    for line in file:
      m1 = re.match(r"^([0-9]{2}:[0-9]{2}:[0-9]{2},[0-9]{3}) --> ([0-9]{2}:[0-9]{2}:[0-9]{2},[0-9]{3})", line)
      if line == "\n":
        if (entryNumber > 0) and (timestamp is not None):
          s = SubtitleEntry(entryNumber, text, timestamp.groups()[0], timestamp.groups()[1])
          Subtitles.append(s)
          text = ""
          timestamp = ""
          entryNumber = 0
      elif re.match(r"^[0-9]+$", line):
        entryNumber = int(line.rstrip())
      elif m1:
        timestamp = m1
      else:
        text += line

  if args.inplace == True:
    args.outfile = args.file
    args.file = args.file + ".backup"
    with open(args.file, 'w') as o:
      print("Writing backup to " + args.file)
      for entry in Subtitles:
        print(entry.formatEntry(), file=o)

  print("Updating subtitle file: " + args.outfile)
  lastEntryNumber = Subtitles[-1].entryNumber
  with open(args.outfile, 'w') as o:
    print("Shifting all entries by {} seconds".format(args.seconds))
    for idx, entry in enumerate(Subtitles):
      if idx == 0:
        print("First subtitle line: {}".format(entry.text))
        if args.strip_first and entry.text != "":
          print("Stripping first line - {}".format(entry.text))
          entry.text = ""
      if idx == lastEntryNumber - 1:
        print("Last subtitle line: {}".format(entry.text))
        if args.strip_last and entry.text != "":
          print("Stripping last line - {}".format(entry.text))
          entry.text = ""
      entry.shift(args.seconds)
      print(entry.formatEntry(), file=o)

if __name__ == "__main__":
  main()