From 4afc69f489cc1b69a32104d8c0c9930b952899bb Mon Sep 17 00:00:00 2001
From: Paxx <daniel.botta@hiway.media>
Date: Tue, 4 Jun 2024 10:33:26 +0200
Subject: [PATCH] release: added helper methods on probe

---
 ffprobe/probe.go | 34 +++++++++++++++++++++++++---------
 1 file changed, 25 insertions(+), 9 deletions(-)

diff --git a/ffprobe/probe.go b/ffprobe/probe.go
index ef216f9..cc9c739 100644
--- a/ffprobe/probe.go
+++ b/ffprobe/probe.go
@@ -93,20 +93,36 @@ type Tags struct {
 	VendorId     string    `json:"vendor_id"`
 }
 
-func (p *Probe) GetFirstVideoStream() *Stream {
+func (p *Probe) GetStreams(streamType string) []Stream {
+	result := make([]Stream, 0)
 	for _, v := range p.Streams {
-		if v.CodecType == "video" {
-			return &v
+		if v.CodecType == streamType {
+			result = append(result, v)
 		}
 	}
 	return nil
 }
 
+func (p *Probe) GetFirstVideoStream() *Stream {
+	return &p.GetStreams("video")[0]
+}
+
 func (p *Probe) GetFirstAudioStream() *Stream {
-	for _, v := range p.Streams {
-		if v.CodecType == "audio" {
-			return &v
-		}
-	}
-	return nil
+	return &p.GetStreams("audio")[0]
+}
+
+func (p *Probe) GetVideoStreams() []Stream {
+	return p.GetStreams("video")
+}
+
+func (p *Probe) GetAudioStreams() []Stream {
+	return p.GetStreams("audio")
+}
+
+func (p *Probe) GetSubtitleStreams() []Stream {
+	return p.GetStreams("subtitle")
+}
+
+func (p *Probe) GetDataStreams() []Stream {
+	return p.GetStreams("data")
 }