Merge pull request #26 from wolveix/add-missing-stream-types

Added two missing data stream constants & added to test

ffprobe_test.go

@@ -84,14 +84,24 @@ func validateData(t *testing.T, data *ProbeData) {
 		t.Errorf("It does not have a subtitle stream.")
 	}
 
+	stream = data.StreamType(StreamData)
+	if len(stream) != 0 {
+		t.Errorf("It does not have a data stream.")
+	}
+
+	stream = data.StreamType(StreamAttachment)
+	if len(stream) != 0 {
+		t.Errorf("It does not have an attachment stream.")
+	}
+
 	stream = data.StreamType(StreamAny)
 	if len(stream) != 2 {
 		t.Errorf("It should have two streams.")
 	}
 
 	// test Format.Duration
-	dration := data.Format.Duration()
-	if dration.Seconds() != 5.312 {
+	duration := data.Format.Duration()
+	if duration.Seconds() != 5.312 {
 		t.Errorf("this video is 5.312s.")
 	}
 	// test Format.StartTime

probedata.go

@@ -16,6 +16,10 @@ const (
 	StreamAudio StreamType = "audio"
 	// StreamSubtitle is a subtitle stream
 	StreamSubtitle StreamType = "subtitle"
+	// StreamData is a data stream
+	StreamData StreamType = "data"
+	// StreamAttachment is an attachment stream
+	StreamAttachment StreamType = "attachment"
 )
 
 // ProbeData is the root json data structure returned by an ffprobe.
@@ -142,39 +146,37 @@ func (p *ProbeData) StreamType(streamType StreamType) (streams []Stream) {
 
 // FirstVideoStream returns the first video stream found
 func (p *ProbeData) FirstVideoStream() *Stream {
-	for _, s := range p.Streams {
-		if s == nil {
-			continue
-		}
-		if s.CodecType == string(StreamVideo) {
-			return s
-		}
-	}
-	return nil
+	return p.firstStream(StreamVideo)
 }
 
 // FirstAudioStream returns the first audio stream found
 func (p *ProbeData) FirstAudioStream() *Stream {
-	for _, s := range p.Streams {
-		if s == nil {
-			continue
-		}
-		if s.CodecType == string(StreamAudio) {
-			return s
-		}
-	}
-	return nil
+	return p.firstStream(StreamAudio)
 }
 
 // FirstSubtitleStream returns the first subtitle stream found
 func (p *ProbeData) FirstSubtitleStream() *Stream {
+	return p.firstStream(StreamSubtitle)
+}
+
+// FirstDataStream returns the first data stream found
+func (p *ProbeData) FirstDataStream() *Stream {
+	return p.firstStream(StreamData)
+}
+
+// FirstAttachmentStream returns the first attachment stream found
+func (p *ProbeData) FirstAttachmentStream() *Stream {
+	return p.firstStream(StreamAttachment)
+}
+
+func (p *ProbeData) firstStream(streamType StreamType) *Stream {
 	for _, s := range p.Streams {
 		if s == nil {
 			continue
 		}
-		if s.CodecType == string(StreamSubtitle) {
+		if s.CodecType == string(streamType) {
 			return s
 		}
 	}
 	return nil
-}
+}