Support SRT styling and bold/underline in VTT

srt.go

@@ -4,9 +4,12 @@ import (
 	"bufio"
 	"fmt"
 	"io"
+	"regexp"
 	"strconv"
 	"strings"
 	"time"
+
+	"golang.org/x/net/html"
 )
 
 // Constants
@@ -17,6 +20,7 @@ const (
 // Vars
 var (
 	bytesSRTTimeBoundariesSeparator = []byte(srtTimeBoundariesSeparator)
+	regexpSRTSSATags                = regexp.MustCompile(`{\\.*?}`)
 )
 
 // parseDurationSRT parses an .srt duration
@@ -111,12 +115,134 @@ 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(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 != "" {
+				// Remove all SSA/ASS tags from text
+				s := regexpSRTSSATags.ReplaceAllStringFunc(s, removeSSATagsWithPos(&pos))
+				// 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
+}
+
+// Removes SSA/ASS tags from subtitle text
+// and extracts position if detected
+func removeSSATagsWithPos(pos *byte) func(string) string {
+	return func(i string) string {
+		// Based on in the following information:
+		// https://superuser.com/a/1228528
+		switch i {
+		case `{\an7}`: // top-left
+			*pos = 7
+		case `{\an8}`: // top-center
+			*pos = 8
+		case `{\an9}`: // top-right
+			*pos = 9
+		case `{\an4}`: // middle-left
+			*pos = 4
+		case `{\an5}`: // middle-center
+			*pos = 5
+		case `{\an6}`: // middle-right
+			*pos = 6
+		case `{\an1}`: // bottom-left
+			*pos = 1
+		case `{\an2}`: // bottom-center
+			*pos = 2
+		case `{\an3}`: // bottom-right
+			*pos = 3
+		}
+
+		// Remove tag from subtitle text
+		return ""
+	}
+}
+
 // formatDurationSRT formats an .srt duration
 func formatDurationSRT(i time.Duration) string {
 	return formatDuration(i, ",", 3)
@@ -146,8 +272,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, l.srtBytes()...)
 		}
 
 		// Add new line
@@ -164,3 +289,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("<font color=\""+color+"\">")...)
+	}
+	if b {
+		c = append(c, []byte("<b>")...)
+	}
+	if i {
+		c = append(c, []byte("<i>")...)
+	}
+	if u {
+		c = append(c, []byte("<u>")...)
+	}
+	if pos != 0 {
+		c = append(c, []byte(fmt.Sprintf(`{\an%d}`, pos))...)
+	}
+	c = append(c, []byte(li.Text)...)
+	if u {
+		c = append(c, []byte("</u>")...)
+	}
+	if i {
+		c = append(c, []byte("</i>")...)
+	}
+	if b {
+		c = append(c, []byte("</b>")...)
+	}
+	if color != "" {
+		c = append(c, []byte("</font>")...)
+	}
+	return
+}

srt_test.go

@@ -46,3 +46,23 @@ func TestSRTMissingSequence(t *testing.T) {
 	assert.NoError(t, err)
 	assert.Equal(t, string(c), w.String())
 }
+
+func TestSRTStyled(t *testing.T) {
+	// Open
+	s, err := astisub.OpenFile("./testdata/example-styled-in.srt")
+	assert.NoError(t, err)
+	assertStyledSubtitleItems(t, s)
+	assertSRTSubtitleStyles(t, s)
+
+	// No subtitles to write
+	w := &bytes.Buffer{}
+	err = astisub.Subtitles{}.WriteToSRT(w)
+	assert.EqualError(t, err, astisub.ErrNoSubtitlesToWrite.Error())
+
+	// Write
+	c, err := ioutil.ReadFile("./testdata/example-styled-out.srt")
+	assert.NoError(t, err)
+	err = s.WriteToSRT(w)
+	assert.NoError(t, err)
+	assert.Equal(t, string(c), w.String())
+}

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,7 +242,9 @@ type StyleAttributes struct {
 	TTMLWritingMode      *string
 	TTMLZIndex           *int
 	WebVTTAlign          string
+	WebVTTBold           bool
 	WebVTTItalics        bool
+	WebVTTUnderline      bool
 	WebVTTLine           string
 	WebVTTLines          int
 	WebVTTPosition       string
@@ -248,6 +256,45 @@ type StyleAttributes struct {
 	WebVTTWidth          string
 }
 
+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
+}
+
 func (sa *StyleAttributes) propagateSSAAttributes() {}
 
 func (sa *StyleAttributes) propagateSTLAttributes() {
@@ -321,7 +368,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
@@ -802,3 +857,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
+}