Support ASS position tags in SRT

asticode · Oct 23, 2023 · d8d439c · d8d439c
1 parent e2f0067
commit d8d439c
Show file tree

Hide file tree

Showing 5 changed files with 109 additions and 24 deletions.
diff --git a/srt.go b/srt.go
@@ -140,6 +140,7 @@ func parseTextSrt(i string) (o Line) {
 		italic    bool
 		underline bool
 		color     *string
+		pos       byte
 	)
 	for {
 		// Get next tag
@@ -184,16 +185,16 @@ func parseTextSrt(i string) (o Line) {
 			}
 		case html.TextToken:
 			if s := strings.TrimSpace(raw); s != "" {
-				// remove all ssa/ass tags from text
-				// TODO: maybe add support for {\an8}
-				s := regexpSRTSSATags.ReplaceAllLiteralString(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 {
+				if bold || italic || underline || color != nil || pos != 0 {
 					sa = &StyleAttributes{
-						SRTColor:     color,
 						SRTBold:      bold,
+						SRTColor:     color,
 						SRTItalics:   italic,
+						SRTPosition:  pos,
 						SRTUnderline: underline,
 					}
 					sa.propagateSRTAttributes()
@@ -210,6 +211,38 @@ func parseTextSrt(i string) (o Line) {
 	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)
@@ -281,6 +314,12 @@ func (li LineItem) srtBytes() (c []byte) {
 	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+"\">")...)
@@ -294,6 +333,9 @@ func (li LineItem) srtBytes() (c []byte) {
 	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>")...)

diff --git a/srt_test.go b/srt_test.go
@@ -52,6 +52,7 @@ func TestSRTStyled(t *testing.T) {
 	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{}

diff --git a/subtitles.go b/subtitles.go
@@ -174,9 +174,10 @@ var (
 
 // StyleAttributes represents style attributes
 type StyleAttributes struct {
-	SRTColor             *string
 	SRTBold              bool
+	SRTColor             *string
 	SRTItalics           bool
+	SRTPosition          byte // 1-9 numpad layout
 	SRTUnderline         bool
 	SSAAlignment         *int
 	SSAAlphaLevel        *float64
@@ -261,6 +262,34 @@ func (sa *StyleAttributes) propagateSRTAttributes() {
 		// 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

diff --git a/subtitles_test.go b/subtitles_test.go
@@ -45,41 +45,54 @@ func assertStyledSubtitleItems(t *testing.T, i *astisub.Subtitles) {
 	assert.Equal(t, 17*time.Second+985*time.Millisecond, i.Items[0].StartAt)
 	assert.Equal(t, 20*time.Second+521*time.Millisecond, i.Items[0].EndAt)
 	assert.Equal(t, "[instrumental music]", i.Items[0].Lines[0].String())
+	assert.Equal(t, 47*time.Second+115*time.Millisecond, i.Items[1].StartAt)
+	assert.Equal(t, 48*time.Second+282*time.Millisecond, i.Items[1].EndAt)
+	assert.Equal(t, "[ticks]", i.Items[1].Lines[0].String())
+	assert.Equal(t, 58*time.Second+192*time.Millisecond, i.Items[2].StartAt)
+	assert.Equal(t, 59*time.Second+727*time.Millisecond, i.Items[2].EndAt)
+	assert.Equal(t, "[instrumental music]", i.Items[2].Lines[0].String())
+	assert.Equal(t, 1*time.Minute+1*time.Second+662*time.Millisecond, i.Items[3].StartAt)
+	assert.Equal(t, 1*time.Minute+3*time.Second+63*time.Millisecond, i.Items[3].EndAt)
+	assert.Equal(t, "[dog barking]", i.Items[3].Lines[0].String())
+	assert.Equal(t, 1*time.Minute+26*time.Second+787*time.Millisecond, i.Items[4].StartAt)
+	assert.Equal(t, 1*time.Minute+29*time.Second+523*time.Millisecond, i.Items[4].EndAt)
+	assert.Equal(t, "[beeping]", i.Items[4].Lines[0].String())
+	assert.Equal(t, 1*time.Minute+29*time.Second+590*time.Millisecond, i.Items[5].StartAt)
+	assert.Equal(t, 1*time.Minute+31*time.Second+992*time.Millisecond, i.Items[5].EndAt)
+	assert.Equal(t, "[automated]", i.Items[5].Lines[0].String())
+	assert.Equal(t, "'The time is 7:35.'", i.Items[5].Lines[1].String())
+}
+
+func assertSRTSubtitleStyles(t *testing.T, i *astisub.Subtitles) {
+	assert.Len(t, i.Items, 6)
 	assert.Equal(t, "#00ff00", *i.Items[0].Lines[0].Items[0].InlineStyle.SRTColor)
+	assert.Zero(t, i.Items[0].Lines[0].Items[0].InlineStyle.SRTPosition)
 	assert.True(t, i.Items[0].Lines[0].Items[0].InlineStyle.SRTBold)
 	assert.False(t, i.Items[0].Lines[0].Items[0].InlineStyle.SRTItalics)
 	assert.False(t, i.Items[0].Lines[0].Items[0].InlineStyle.SRTUnderline)
-	assert.Equal(t, 47*time.Second+115*time.Millisecond, i.Items[1].StartAt)
-	assert.Equal(t, 48*time.Second+282*time.Millisecond, i.Items[1].EndAt)
-	assert.Equal(t, "[ticks]", i.Items[1].Lines[0].String())
 	assert.Equal(t, "#ffffff", *i.Items[1].Lines[0].Items[0].InlineStyle.SRTColor)
+	assert.Zero(t, i.Items[1].Lines[0].Items[0].InlineStyle.SRTPosition)
 	assert.False(t, i.Items[1].Lines[0].Items[0].InlineStyle.SRTBold)
 	assert.False(t, i.Items[1].Lines[0].Items[0].InlineStyle.SRTItalics)
 	assert.False(t, i.Items[1].Lines[0].Items[0].InlineStyle.SRTUnderline)
-	assert.Equal(t, 58*time.Second+192*time.Millisecond, i.Items[2].StartAt)
-	assert.Equal(t, 59*time.Second+727*time.Millisecond, i.Items[2].EndAt)
-	assert.Equal(t, "[instrumental music]", i.Items[2].Lines[0].String())
 	assert.Equal(t, "#00ff00", *i.Items[2].Lines[0].Items[0].InlineStyle.SRTColor)
+	assert.Zero(t, i.Items[2].Lines[0].Items[0].InlineStyle.SRTPosition)
 	assert.False(t, i.Items[2].Lines[0].Items[0].InlineStyle.SRTBold)
 	assert.False(t, i.Items[2].Lines[0].Items[0].InlineStyle.SRTItalics)
 	assert.False(t, i.Items[2].Lines[0].Items[0].InlineStyle.SRTUnderline)
-	assert.Equal(t, 1*time.Minute+1*time.Second+662*time.Millisecond, i.Items[3].StartAt)
-	assert.Equal(t, 1*time.Minute+3*time.Second+63*time.Millisecond, i.Items[3].EndAt)
-	assert.Equal(t, "[dog barking]", i.Items[3].Lines[0].String())
 	assert.Nil(t, i.Items[3].Lines[0].Items[0].InlineStyle.SRTColor)
+	assert.Zero(t, i.Items[3].Lines[0].Items[0].InlineStyle.SRTPosition)
 	assert.True(t, i.Items[3].Lines[0].Items[0].InlineStyle.SRTBold)
 	assert.False(t, i.Items[3].Lines[0].Items[0].InlineStyle.SRTItalics)
 	assert.True(t, i.Items[3].Lines[0].Items[0].InlineStyle.SRTUnderline)
-	assert.Equal(t, 1*time.Minute+26*time.Second+787*time.Millisecond, i.Items[4].StartAt)
-	assert.Equal(t, 1*time.Minute+29*time.Second+523*time.Millisecond, i.Items[4].EndAt)
-	assert.Equal(t, "[beeping]", i.Items[4].Lines[0].String())
-	assert.Nil(t, i.Items[4].Lines[0].Items[0].InlineStyle)
-	assert.Equal(t, 1*time.Minute+29*time.Second+590*time.Millisecond, i.Items[5].StartAt)
-	assert.Equal(t, 1*time.Minute+31*time.Second+992*time.Millisecond, i.Items[5].EndAt)
-	assert.Equal(t, "[automated]", i.Items[5].Lines[0].String())
+	assert.Nil(t, i.Items[4].Lines[0].Items[0].InlineStyle.SRTColor)
+	assert.Equal(t, byte(8), i.Items[4].Lines[0].Items[0].InlineStyle.SRTPosition)
+	assert.False(t, i.Items[4].Lines[0].Items[0].InlineStyle.SRTBold)
+	assert.False(t, i.Items[4].Lines[0].Items[0].InlineStyle.SRTItalics)
+	assert.False(t, i.Items[4].Lines[0].Items[0].InlineStyle.SRTUnderline)
 	assert.Nil(t, i.Items[5].Lines[0].Items[0].InlineStyle)
-	assert.Equal(t, "'The time is 7:35.'", i.Items[5].Lines[1].String())
 	assert.Nil(t, i.Items[5].Lines[1].Items[0].InlineStyle.SRTColor)
+	assert.Zero(t, i.Items[5].Lines[1].Items[0].InlineStyle.SRTPosition)
 	assert.False(t, i.Items[5].Lines[1].Items[0].InlineStyle.SRTBold)
 	assert.True(t, i.Items[5].Lines[1].Items[0].InlineStyle.SRTItalics)
 	assert.False(t, i.Items[5].Lines[1].Items[0].InlineStyle.SRTUnderline)

diff --git a/testdata/example-styled-out.srt b/testdata/example-styled-out.srt
@@ -16,7 +16,7 @@
 
 5
 00:01:26,787 --> 00:01:29,523
-[beeping]
+{\an8}[beeping]
 
 6
 00:01:29,590 --> 00:01:31,992