#!/usr/bin/env python

'''A tiny library for parsing, modifying, and composing SRT files.'''

import functools
import re
from datetime import timedelta
from itertools import groupby


SUBTITLE_PATTERN = r'(\d+)\n(\d+:\d+:\d+,\d+) --> (\d+:\d+:\d+,\d+)\n(.+?)\n\n'
SUBTITLE_REGEX = re.compile(SUBTITLE_PATTERN, re.MULTILINE | re.DOTALL)


@functools.total_ordering
class Subtitle(object):
    def __init__(self, index, start, end, content):
        self.index = index
        self.start = start
        self.end = end
        self.content = content

    def __eq__(self, other):
        return self.__dict__ == other.__dict__

    def __lt__(self, other):
        return self.start < other.start

    def to_srt(self):
        return '%d\n%s --> %s\n%s\n\n' % (
            self.index, timedelta_to_srt_timestamp(self.start),
            timedelta_to_srt_timestamp(self.end), self.content,
        )


def timedelta_to_srt_timestamp(timedelta_timestamp):
    '''Convert a timedelta to an SRT timestamp.'''
    hrs, remainder = divmod(timedelta_timestamp.seconds, 3600)
    mins, secs = divmod(remainder, 60)
    msecs = timedelta_timestamp.microseconds // 1000
    return '%02d:%02d:%02d,%03d' % (hrs, mins, secs, msecs)


def srt_timestamp_to_timedelta(srt_timestamp):
    '''Convert an SRT timestamp to a timedelta.'''
    hrs, mins, secs, msecs = map(int, re.split('[,:]', srt_timestamp))
    return timedelta(hours=hrs, minutes=mins, seconds=secs, milliseconds=msecs)


def parse(srt):
    '''Convert an SRT formatted string to a generator of Subtitle objects.'''
    for match in SUBTITLE_REGEX.finditer(srt):
        raw_index, raw_start, raw_end, content = match.groups()
        yield Subtitle(
            index=int(raw_index), start=srt_timestamp_to_timedelta(raw_start),
            end=srt_timestamp_to_timedelta(raw_end), content=content,
        )


def parse_file(srt_stream):
    '''Parse an SRT formatted stream into Subtitle objects.'''
    for is_sep, lines in groupby(srt_stream, lambda line: line != '\n'):
        if is_sep:
            srt_block = ''.join(lines) + '\n'
            subtitle, = parse(srt_block)
            yield subtitle


def compose(subtitles):
    '''Convert an iterator of Subtitle objects to SRT formatted strings.'''
    return ''.join(subtitle.to_srt() for subtitle in subtitles)


def compose_file(subtitles, srt_stream):
    '''Stream a sequence of Subtitle objects into an SRT formatted stream.'''
    for subtitle in subtitles:
        srt_stream.write(subtitle.to_srt())

@@ -377,24 +377,79 @@
 le(object):%0A
+    '''The metadata relating to a single SRT block.'''%0A
     def __in
@@ -766,16 +766,73 @@
 (self):%0A
+        '''Convert a Subtitle object to an SRT block.'''%0A
         

Add docstrings to Subtitle class