diff --git a/srt.go b/srt.go index 58f77ce..7b8d0c9 100644 --- a/srt.go +++ b/srt.go @@ -8,6 +8,8 @@ import ( "strings" "time" "unicode/utf8" + + "golang.org/x/net/html" ) // Constants @@ -116,7 +118,95 @@ func ReadFromSRT(i io.Reader) (o *Subtitles, err error) { o.Items = append(o.Items, s) } else { // Add text - s.Lines = append(s.Lines, Line{Items: []LineItem{{Text: strings.TrimSpace(line)}}}) + if l := parseTextSrt(strings.TrimSpace(line)); len(l.Items) > 0 { + s.Lines = append(s.Lines, l) + } + } + } + return +} + +// parseTextSrt parses the input line to fill the Line +func parseTextSrt(i string) (o Line) { + // special handling needed for empty line + if strings.TrimSpace(i) == "" { + o.Items = []LineItem{{Text: ""}} + return + } + + // Create tokenizer + tr := html.NewTokenizer(strings.NewReader(i)) + + // Loop + var ( + bold bool + italic bool + underline bool + color *string + pos byte + ) + for { + // Get next tag + t := tr.Next() + + // Process error + if err := tr.Err(); err != nil { + break + } + + // Get unmodified text + raw := string(tr.Raw()) + // Get current token + token := tr.Token() + + switch t { + case html.EndTagToken: + // Parse italic/bold/underline + switch token.Data { + case "b": + bold = false + case "i": + italic = false + case "u": + underline = false + case "font": + color = nil + } + case html.StartTagToken: + // Parse italic/bold/underline + switch token.Data { + case "b": + bold = true + case "i": + italic = true + case "u": + underline = true + case "font": + if c := htmlTokenAttribute(&token, "color"); c != nil { + color = c + } + } + case html.TextToken: + if s := strings.TrimSpace(raw); s != "" { + // Get style attribute + var sa *StyleAttributes + if bold || italic || underline || color != nil || pos != 0 { + sa = &StyleAttributes{ + SRTBold: bold, + SRTColor: color, + SRTItalics: italic, + SRTPosition: pos, + SRTUnderline: underline, + } + sa.propagateSRTAttributes() + } + + // Append item + o.Items = append(o.Items, LineItem{ + InlineStyle: sa, + Text: s, + }) + } } } return @@ -151,8 +241,7 @@ func (s Subtitles) WriteToSRT(o io.Writer) (err error) { // Loop through lines for _, l := range v.Lines { - c = append(c, []byte(l.String())...) - c = append(c, bytesLineSeparator...) + c = append(c, []byte(l.srtBytes())...) } // Add new line @@ -169,3 +258,65 @@ func (s Subtitles) WriteToSRT(o io.Writer) (err error) { } return } + +func (l Line) srtBytes() (c []byte) { + for idx, li := range l.Items { + c = append(c, li.srtBytes()...) + // condition to avoid adding space as the last character. + if idx < len(l.Items)-1 { + c = append(c, []byte(" ")...) + } + } + c = append(c, bytesLineSeparator...) + return +} + +func (li LineItem) srtBytes() (c []byte) { + // Get color + var color string + if li.InlineStyle != nil && li.InlineStyle.SRTColor != nil { + color = *li.InlineStyle.SRTColor + } + + // Get bold/italics/underline + b := li.InlineStyle != nil && li.InlineStyle.SRTBold + i := li.InlineStyle != nil && li.InlineStyle.SRTItalics + u := li.InlineStyle != nil && li.InlineStyle.SRTUnderline + + // Get position + var pos byte + if li.InlineStyle != nil { + pos = li.InlineStyle.SRTPosition + } + + // Append + if color != "" { + c = append(c, []byte("")...) + } + if b { + c = append(c, []byte("")...) + } + if i { + c = append(c, []byte("")...) + } + if u { + c = append(c, []byte("")...) + } + if pos != 0 { + c = append(c, []byte(fmt.Sprintf(`{\an%d}`, pos))...) + } + c = append(c, []byte(li.Text)...) + if u { + c = append(c, []byte("")...) + } + if i { + c = append(c, []byte("")...) + } + if b { + c = append(c, []byte("")...) + } + if color != "" { + c = append(c, []byte("")...) + } + return +} diff --git a/srt_test.go b/srt_test.go index 611c3fc..3ccb888 100644 --- a/srt_test.go +++ b/srt_test.go @@ -3,7 +3,9 @@ package astisub_test import ( "bytes" "io/ioutil" + "os" "testing" + "time" "github.com/asticode/go-astisub" "github.com/stretchr/testify/assert" @@ -51,3 +53,77 @@ func TestNonUTF8SRT(t *testing.T) { _, err := astisub.OpenFile("./testdata/example-in-non-utf8.srt") assert.Error(t, err) } + +func TestSRTStyled(t *testing.T) { + // Open + s, err := astisub.OpenFile("./testdata/example-in-styled.srt") + assert.NoError(t, err) + + // assert the items are properly parsed + assert.Len(t, s.Items, 6) + assert.Equal(t, 17*time.Second+985*time.Millisecond, s.Items[0].StartAt) + assert.Equal(t, 20*time.Second+521*time.Millisecond, s.Items[0].EndAt) + assert.Equal(t, "[instrumental music]", s.Items[0].Lines[0].String()) + assert.Equal(t, 47*time.Second+115*time.Millisecond, s.Items[1].StartAt) + assert.Equal(t, 48*time.Second+282*time.Millisecond, s.Items[1].EndAt) + assert.Equal(t, "[ticks]", s.Items[1].Lines[0].String()) + assert.Equal(t, 58*time.Second+192*time.Millisecond, s.Items[2].StartAt) + assert.Equal(t, 59*time.Second+727*time.Millisecond, s.Items[2].EndAt) + assert.Equal(t, "[instrumental music]", s.Items[2].Lines[0].String()) + assert.Equal(t, 1*time.Minute+1*time.Second+662*time.Millisecond, s.Items[3].StartAt) + assert.Equal(t, 1*time.Minute+3*time.Second+63*time.Millisecond, s.Items[3].EndAt) + assert.Equal(t, "[dog barking]", s.Items[3].Lines[0].String()) + assert.Equal(t, 1*time.Minute+26*time.Second+787*time.Millisecond, s.Items[4].StartAt) + assert.Equal(t, 1*time.Minute+29*time.Second+523*time.Millisecond, s.Items[4].EndAt) + assert.Equal(t, "[beeping]", s.Items[4].Lines[0].String()) + assert.Equal(t, 1*time.Minute+29*time.Second+590*time.Millisecond, s.Items[5].StartAt) + assert.Equal(t, 1*time.Minute+31*time.Second+992*time.Millisecond, s.Items[5].EndAt) + assert.Equal(t, "[automated]", s.Items[5].Lines[0].String()) + assert.Equal(t, "'The time is 7:35.'", s.Items[5].Lines[1].String()) + + // assert the styles of the items + assert.Len(t, s.Items, 6) + assert.Equal(t, "#00ff00", *s.Items[0].Lines[0].Items[0].InlineStyle.SRTColor) + assert.Zero(t, s.Items[0].Lines[0].Items[0].InlineStyle.SRTPosition) + assert.True(t, s.Items[0].Lines[0].Items[0].InlineStyle.SRTBold) + assert.False(t, s.Items[0].Lines[0].Items[0].InlineStyle.SRTItalics) + assert.False(t, s.Items[0].Lines[0].Items[0].InlineStyle.SRTUnderline) + assert.Equal(t, "#ff00ff", *s.Items[1].Lines[0].Items[0].InlineStyle.SRTColor) + assert.Zero(t, s.Items[1].Lines[0].Items[0].InlineStyle.SRTPosition) + assert.False(t, s.Items[1].Lines[0].Items[0].InlineStyle.SRTBold) + assert.False(t, s.Items[1].Lines[0].Items[0].InlineStyle.SRTItalics) + assert.False(t, s.Items[1].Lines[0].Items[0].InlineStyle.SRTUnderline) + assert.Equal(t, "#00ff00", *s.Items[2].Lines[0].Items[0].InlineStyle.SRTColor) + assert.Zero(t, s.Items[2].Lines[0].Items[0].InlineStyle.SRTPosition) + assert.False(t, s.Items[2].Lines[0].Items[0].InlineStyle.SRTBold) + assert.False(t, s.Items[2].Lines[0].Items[0].InlineStyle.SRTItalics) + assert.False(t, s.Items[2].Lines[0].Items[0].InlineStyle.SRTUnderline) + assert.Nil(t, s.Items[3].Lines[0].Items[0].InlineStyle.SRTColor) + assert.Zero(t, s.Items[3].Lines[0].Items[0].InlineStyle.SRTPosition) + assert.True(t, s.Items[3].Lines[0].Items[0].InlineStyle.SRTBold) + assert.False(t, s.Items[3].Lines[0].Items[0].InlineStyle.SRTItalics) + assert.True(t, s.Items[3].Lines[0].Items[0].InlineStyle.SRTUnderline) + assert.Nil(t, s.Items[4].Lines[0].Items[0].InlineStyle) + assert.Nil(t, s.Items[5].Lines[0].Items[0].InlineStyle) + assert.Nil(t, s.Items[5].Lines[1].Items[0].InlineStyle.SRTColor) + assert.Zero(t, s.Items[5].Lines[1].Items[0].InlineStyle.SRTPosition) + assert.False(t, s.Items[5].Lines[1].Items[0].InlineStyle.SRTBold) + assert.True(t, s.Items[5].Lines[1].Items[0].InlineStyle.SRTItalics) + assert.False(t, s.Items[5].Lines[1].Items[0].InlineStyle.SRTUnderline) + + // Write to srt + w := &bytes.Buffer{} + c, err := os.ReadFile("./testdata/example-out-styled.srt") + assert.NoError(t, err) + err = s.WriteToSRT(w) + assert.NoError(t, err) + assert.Equal(t, string(c), w.String()) + + // Write to WebVTT + w = &bytes.Buffer{} + c, err = os.ReadFile("./testdata/example-out-styled.vtt") + assert.NoError(t, err) + err = s.WriteToWebVTT(w) + assert.NoError(t, err) + assert.Equal(t, string(c), w.String()) +} diff --git a/subtitles.go b/subtitles.go index e917931..ff4811f 100644 --- a/subtitles.go +++ b/subtitles.go @@ -12,6 +12,7 @@ import ( "time" "github.com/asticode/go-astikit" + "golang.org/x/net/html" ) // Bytes @@ -173,6 +174,11 @@ var ( // StyleAttributes represents style attributes type StyleAttributes struct { + SRTBold bool + SRTColor *string + SRTItalics bool + SRTPosition byte // 1-9 numpad layout + SRTUnderline bool SSAAlignment *int SSAAlphaLevel *float64 SSAAngle *float64 // degrees @@ -236,6 +242,8 @@ type StyleAttributes struct { TTMLWritingMode *string TTMLZIndex *int WebVTTAlign string + WebVTTBold bool + WebVTTItalics bool WebVTTLine string WebVTTLines int WebVTTPosition string @@ -244,6 +252,7 @@ type StyleAttributes struct { WebVTTSize string WebVTTStyles []string WebVTTTags []WebVTTTag + WebVTTUnderline bool WebVTTVertical string WebVTTViewportAnchor string WebVTTWidth string @@ -279,6 +288,56 @@ func (t WebVTTTag) endTag() string { return "" } +func (sa *StyleAttributes) propagateSRTAttributes() { + // copy relevant attrs to WebVTT ones + if sa.SRTColor != nil { + // TODO: handle non-default colors that need custom styles + sa.TTMLColor = sa.SRTColor + } + + switch sa.SRTPosition { + case 7: // top-left + sa.WebVTTAlign = "left" + sa.WebVTTPosition = "10%" + case 8: // top-center + sa.WebVTTPosition = "10%" + case 9: // top-right + sa.WebVTTAlign = "right" + sa.WebVTTPosition = "10%" + case 4: // middle-left + sa.WebVTTAlign = "left" + sa.WebVTTPosition = "50%" + case 5: // middle-center + sa.WebVTTPosition = "50%" + case 6: // middle-right + sa.WebVTTAlign = "right" + sa.WebVTTPosition = "50%" + case 1: // bottom-left + sa.WebVTTAlign = "left" + sa.WebVTTPosition = "90%" + case 2: // bottom-center + sa.WebVTTPosition = "90%" + case 3: // bottom-right + sa.WebVTTAlign = "right" + sa.WebVTTPosition = "90%" + } + + sa.WebVTTBold = sa.SRTBold + sa.WebVTTItalics = sa.SRTItalics + sa.WebVTTUnderline = sa.SRTUnderline + + sa.WebVTTTags = make([]WebVTTTag, 0) + if sa.WebVTTBold { + sa.WebVTTTags = append(sa.WebVTTTags, WebVTTTag{Name: "b"}) + } + if sa.WebVTTItalics { + sa.WebVTTTags = append(sa.WebVTTTags, WebVTTTag{Name: "i"}) + } + if sa.WebVTTUnderline { + sa.WebVTTTags = append(sa.WebVTTTags, WebVTTTag{Name: "u"}) + } +} + func (sa *StyleAttributes) propagateSSAAttributes() {} func (sa *StyleAttributes) propagateSTLAttributes() { @@ -352,7 +411,15 @@ func (sa *StyleAttributes) propagateTTMLAttributes() { } } -func (sa *StyleAttributes) propagateWebVTTAttributes() {} +func (sa *StyleAttributes) propagateWebVTTAttributes() { + // copy relevant attrs to SRT ones + if sa.TTMLColor != nil { + sa.SRTColor = sa.TTMLColor + } + sa.SRTBold = sa.WebVTTBold + sa.SRTItalics = sa.WebVTTItalics + sa.SRTUnderline = sa.WebVTTUnderline +} // Metadata represents metadata // TODO Merge attributes @@ -835,3 +902,14 @@ func appendStringToBytesWithNewLine(i []byte, s string) (o []byte) { o = append(o, bytesLineSeparator...) return } + +func htmlTokenAttribute(t *html.Token, key string) *string { + + for _, attr := range t.Attr { + if attr.Key == key { + return &attr.Val + } + } + + return nil +} diff --git a/testdata/example-in-styled.srt b/testdata/example-in-styled.srt new file mode 100644 index 0000000..28d5c90 --- /dev/null +++ b/testdata/example-in-styled.srt @@ -0,0 +1,24 @@ +1 +00:00:17,985 --> 00:00:20,521 +[instrumental music] + +2 +00:00:47,115 --> 00:00:48,282 +[ticks] + +3 +00:00:58,192 --> 00:00:59,727 +[instrumental music] + +4 +00:01:01,662 --> 00:01:03,063 +[dog barking] + +5 +00:01:26,787 --> 00:01:29,523 +[beeping] + +6 +00:01:29,590 --> 00:01:31,992 +[automated] +'The time is 7:35.' diff --git a/testdata/example-out-styled.srt b/testdata/example-out-styled.srt new file mode 100644 index 0000000..4f14974 --- /dev/null +++ b/testdata/example-out-styled.srt @@ -0,0 +1,24 @@ +1 +00:00:17,985 --> 00:00:20,521 +[instrumental music] + +2 +00:00:47,115 --> 00:00:48,282 +[ticks] + +3 +00:00:58,192 --> 00:00:59,727 +[instrumental music] + +4 +00:01:01,662 --> 00:01:03,063 +[dog barking] + +5 +00:01:26,787 --> 00:01:29,523 +[beeping] + +6 +00:01:29,590 --> 00:01:31,992 +[automated] +'The time is 7:35.' diff --git a/testdata/example-out-styled.vtt b/testdata/example-out-styled.vtt new file mode 100644 index 0000000..7bd4bb0 --- /dev/null +++ b/testdata/example-out-styled.vtt @@ -0,0 +1,26 @@ +WEBVTT + +1 +00:00:17.985 --> 00:00:20.521 +[instrumental music] + +2 +00:00:47.115 --> 00:00:48.282 +[ticks] + +3 +00:00:58.192 --> 00:00:59.727 +[instrumental music] + +4 +00:01:01.662 --> 00:01:03.063 +[dog barking] + +5 +00:01:26.787 --> 00:01:29.523 +[beeping] + +6 +00:01:29.590 --> 00:01:31.992 +[automated] +'The time is 7:35.'