Merge pull request #1 from abema/add-decode-playlist

Add DecodePlaylist function and Playlist interface

examples/decode/decode.go

@@ -0,0 +1,62 @@
+package main
+
+import (
+	"fmt"
+	"strings"
+
+	m3u8 "github.com/abema/go-simple-m3u8"
+)
+
+const sampleData = `#EXTM3U
+#EXT-X-STREAM-INF:BANDWIDTH=1280000,AVERAGE-BANDWIDTH=1000000
+http://example.com/low.m3u8
+#EXT-X-STREAM-INF:BANDWIDTH=2560000,AVERAGE-BANDWIDTH=2000000
+http://example.com/mid.m3u8
+#EXT-X-STREAM-INF:BANDWIDTH=7680000,AVERAGE-BANDWIDTH=6000000
+http://example.com/hi.m3u8
+#EXT-X-STREAM-INF:BANDWIDTH=65000,CODECS="mp4a.40.5"
+http://example.com/audio-only.m3u8
+`
+
+func main() {
+	playlist, err := m3u8.DecodePlaylist(strings.NewReader(sampleData))
+	if err != nil {
+		panic(err)
+	}
+	fmt.Println("Type:", playlist.Type())
+	fmt.Println("Tags:", len(playlist.Master().Tags))
+	for name, values := range playlist.Master().Tags {
+		fmt.Printf("  %s: %d\n", name, len(values))
+	}
+	fmt.Println("Streams:")
+	for i, stream := range playlist.Master().Streams {
+		fmt.Printf("  %d:\n", i)
+		height, width, _ := stream.Attributes.Resolution()
+		fmt.Println("    Height:", height)
+		fmt.Println("    Width:", width)
+		fmt.Println("    URI:", stream.URI)
+	}
+}
+
+/* Output:
+Type: master
+Tags: 1
+  EXTM3U: 1
+Streams:
+  0:
+    Height: 0
+    Width: 0
+    URI: http://example.com/low.m3u8
+  1:
+    Height: 0
+    Width: 0
+    URI: http://example.com/mid.m3u8
+  2:
+    Height: 0
+    Width: 0
+    URI: http://example.com/hi.m3u8
+  3:
+    Height: 0
+    Width: 0
+    URI: http://example.com/audio-only.m3u8
+*/

master_playlist.go

@@ -133,3 +133,20 @@ func (playlist *MasterPlaylist) Encode(w io.Writer) error {
 	}
 	return nil
 }
+
+// Type returns the type of the playlist.
+func (playlist *MasterPlaylist) Type() PlaylistType {
+	return PlaylistTypeMaster
+}
+
+// Master returns the master playlist.
+// If the playlist is not a master playlist, it returns nil.
+func (playlist *MasterPlaylist) Master() *MasterPlaylist {
+	return playlist
+}
+
+// Media returns the media playlist.
+// If the playlist is not a media playlist, it returns nil.
+func (playlist *MasterPlaylist) Media() *MediaPlaylist {
+	return nil
+}

media_playlist.go

@@ -121,3 +121,20 @@ func (playlist *MediaPlaylist) Encode(w io.Writer) error {
 	}
 	return nil
 }
+
+// Type returns the type of the playlist.
+func (playlist *MediaPlaylist) Type() PlaylistType {
+	return PlaylistTypeMedia
+}
+
+// Master returns the master playlist.
+// If the playlist is not a master playlist, it returns nil.
+func (playlist *MediaPlaylist) Master() *MasterPlaylist {
+	return nil
+}
+
+// Media returns the media playlist.
+// If the playlist is not a media playlist, it returns nil.
+func (playlist *MediaPlaylist) Media() *MediaPlaylist {
+	return playlist
+}

playlist.go

@@ -0,0 +1,68 @@
+package m3u8
+
+import (
+	"bufio"
+	"bytes"
+	"io"
+)
+
+// PlaylistType represents the type of playlist.
+type PlaylistType string
+
+const (
+	// PlaylistTypeMaster represents a master playlist.
+	PlaylistTypeMaster PlaylistType = "master"
+	// PlaylistTypeMedia represents a media playlist.
+	PlaylistTypeMedia PlaylistType = "media"
+)
+
+type Playlist interface {
+	// Encode encodes the playlist to io.Writer.
+	Encode(w io.Writer) error
+
+	// Type returns the type of the playlist.
+	Type() PlaylistType
+
+	// Master returns the master playlist.
+	// If the playlist is not a master playlist, it returns nil.
+	Master() *MasterPlaylist
+
+	// Media returns the media playlist.
+	// If the playlist is not a media playlist, it returns nil.
+	Media() *MediaPlaylist
+}
+
+// DecodePlaylist detects the type of playlist and decodes it from io.Reader.
+func DecodePlaylist(r io.Reader) (Playlist, error) {
+	data, err := io.ReadAll(r)
+	if err != nil {
+		return nil, err
+	}
+	br := bytes.NewReader(data)
+
+	var masterPlaylistTagCount int
+	var mediaPlaylistTagCount int
+	scanner := bufio.NewScanner(br)
+	for scanner.Scan() {
+		line := scanner.Text()
+		if line == "" {
+			continue
+		}
+		tagName := TagName(line)
+		if isMasterPlaylistTag(tagName) {
+			masterPlaylistTagCount++
+		} else if isMediaPlaylistTag(tagName) || IsSegmentTagName(tagName) {
+			mediaPlaylistTagCount++
+		}
+	}
+	if err := scanner.Err(); err != nil {
+		return nil, err
+	}
+
+	br.Seek(0, io.SeekStart)
+	if masterPlaylistTagCount >= mediaPlaylistTagCount {
+		return DecodeMasterPlaylist(br)
+	} else {
+		return DecodeMediaPlaylist(br)
+	}
+}

playlist_test.go

@@ -0,0 +1,35 @@
+package m3u8
+
+import (
+	"bytes"
+	"fmt"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestDecodePlaylist(t *testing.T) {
+	for idx, testData := range []struct {
+		input        string
+		output       string
+		playlistType PlaylistType
+	}{
+		{input: sampleMaster01Input, output: sampleMaster01Output, playlistType: PlaylistTypeMaster},
+		{input: sampleMaster02Input, output: sampleMaster02Output, playlistType: PlaylistTypeMaster},
+		{input: sampleMedia01Input, output: sampleMedia01Output, playlistType: PlaylistTypeMedia},
+		{input: sampleMedia02Input, output: sampleMedia02Output, playlistType: PlaylistTypeMedia},
+	} {
+		t.Run(fmt.Sprintf("%d", idx), func(t *testing.T) {
+			r := bytes.NewReader([]byte(testData.input))
+			playlist, err := DecodePlaylist(r)
+			require.NoError(t, err)
+			require.Equal(t, playlist.Type(), testData.playlistType)
+			require.Equal(t, playlist.Master() != nil, testData.playlistType == PlaylistTypeMaster)
+			require.Equal(t, playlist.Media() != nil, testData.playlistType == PlaylistTypeMedia)
+			w := bytes.NewBuffer(nil)
+			require.NoError(t, playlist.Encode(w))
+			assert.Equal(t, testData.output, w.String())
+		})
+	}
+}

tags.go

@@ -102,6 +102,23 @@ func getTagOrder(name string) int {
 	return math.MaxInt
 }
 
+var masterPlaylistTagSet = map[string]struct{}{
+	TagExtXMedia:           {},
+	TagExtXStreamInf:       {},
+	TagExtXIFrameStreamInf: {},
+	TagExtXSessionData:     {},
+	TagExtXSessionKey:      {},
+}
+
+var mediaPlaylistTagSet = map[string]struct{}{
+	TagExtXTargetDuration:        {},
+	TagExtXPlaylistType:          {},
+	TagExtXIFramesOnly:           {},
+	TagExtXMediaSequence:         {},
+	TagExtXDiscontinuitySequence: {},
+	TagExtXEndlist:               {},
+}
+
 var segmentTagSet = map[string]struct{}{
 	TagExtInf:              {},
 	TagExtXByteRange:       {},
@@ -143,10 +160,20 @@ func AttributeString(line string) string {
 	return line[idx+1:]
 }
 
+func isMasterPlaylistTag(name string) bool {
+	_, ok := masterPlaylistTagSet[name]
+	return ok
+}
+
+func isMediaPlaylistTag(name string) bool {
+	_, ok := mediaPlaylistTagSet[name]
+	return ok
+}
+
 // IsSegmentTagName returns true if the name is a segment tag name.
 func IsSegmentTagName(name string) bool {
-	_, t := segmentTagSet[name]
-	return t
+	_, ok := segmentTagSet[name]
+	return ok
 }
 
 // Attributes represents a set of attributes of a tag.